diff --git a/CLAUDE.md b/CLAUDE.md index a2405e0..b64929e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -24,6 +24,13 @@ tokeye run "shots/*.npy" --output-dir results # batch inference tokeye download big_tf_unet # pre-fetch weights, print cache path tokeye example # write a synthetic demo signal +# Mode-analysis suite +tokeye modespec modes.yaml # classic Mirnov mode-number analysis (vendored pymodespec) +tokeye elmspec "shots/*.npy" # ELM events from the transient channel +tokeye alfvenspec "shots/*.npy" # Alfvén-eigenmode boxes/masks (ae_tf_maskrcnn) +tokeye eigspec # interactive modal ID / SSI (vendored eigspec port) +tokeye modesearch # mode database — design stage, prints the plan + # Lint uv run ruff check . @@ -39,9 +46,11 @@ Source code lives in `src/tokeye/` (installed as `tokeye` package via `uv_build` ### Models (`models/`) Three model families, each with a `model_*.py` and `config_*.py`: -- **big_tf_unet** — primary transformer U-Net for spectrogram segmentation -- **ae_tf_maskrcnn** — alternative Mask R-CNN approach -- **ae_tf_boxrcnn** — alternative Box R-CNN approach +- **big_tf_unet** — primary transformer U-Net for spectrogram segmentation (HF: `nc1/big_tf_unet`) +- **ae_tf_maskrcnn** — Mask R-CNN instance detector, used by `tokeye alfvenspec` (HF: `nc1/ae_tf_maskrcnn`) +- **ae_tf_boxrcnn** — alternative Box R-CNN approach (not registered, no weights) + +`hub.MODEL_REGISTRY` order is load-bearing: `_build_from_state_dict` probes specs in insertion order, so `big_tf_unet` must stay first. `ModelSpec.repo_id` overrides `DEFAULT_REPO_ID` per model. Shared building blocks in `models/modules/`: `unet.py` (base U-Net), `nn.py` (layers), `bsn.py` (boundary segmentation network). @@ -53,7 +62,15 @@ Gradio web interface launched via `tokeye app` (console script) or `python -m to - **Annotate** (`app/tabs/annotate.py`) — manual labeling interface - **Utilities** (`app/tabs/utilities.py`) — miscellaneous tools -Shared core modules (used by both the app and the `tokeye` CLI) live directly under `src/tokeye/`: `hub.py` (model registry + Hugging Face auto-download), `transforms.py` (STFT), `inference.py` (model inference), `api.py` (the `TokEye` class — public Python API, lazily exported from the package root), `batch.py` (headless batch runner), `examples.py` (synthetic demo signal), `cli.py` (the `tokeye` console entry point). +Shared core modules (used by both the app and the `tokeye` CLI) live directly under `src/tokeye/`: `hub.py` (model registry + Hugging Face auto-download), `transforms.py` (STFT), `inference.py` (model inference, U-Net contract), `api.py` (the `TokEye` class — public Python API, lazily exported from the package root), `batch.py` (headless batch runner), `examples.py` (synthetic demo signal), `cli/` (the `tokeye` console entry point — one module per subcommand, heavy imports deferred into `_handle` functions). + +### Mode-analysis suite +- `modespec/classic/` — **vendored** pymodespec (classic Mirnov mode-number analysis); `modespec/deep/` reserves the next-gen single-chord engine (sibling `integratedmode` project). Vendored code policy: minimal-touch, style rules relaxed in `ruff.toml`, every local change listed in the directory's `PROVENANCE.md`. +- `elmspec/` — ELM event extraction from the transient channel (`events.py` is pure numpy, model plumbing in the CLI handler). +- `alfvenspec/` — R-CNN detection wrapper (`inference.py`; list-of-images contract, windowed processing for wide spectrograms). +- `eigspec/` — **vendored** eigspec MATLAB-toolbox port (modal ID, SSI, random projection); sklearn-dependent clustering behind the `eigspec` extra. +- `modesearch/` — design-stage scaffold only (mode database vision). +- Suite roadmap and future ideas: `docs/ROADMAP.md`. ### Training (`training/`) Multi-step data pipelines (step_0 through step_7) for preparing training data from raw signals. Two regimes: `big_tf_unet/` (original) and `big_tf_unet_multiscale/` (enhanced). Uses PyTorch Lightning. diff --git a/README.md b/README.md index e6218d0..f6b0aac 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,20 @@ tokeye download big_tf_unet # on the login node; prints the cached path tokeye run ... --model big_tf_unet # on the compute node — model is already cached ``` +## Mode-analysis suite + +Beyond segmentation, `tokeye` bundles the analyses DIII-D researchers usually reach for separate tools to get. Each is a subcommand; `--help` on any of them shows the full flags. + +| Command | What it does | +| --- | --- | +| `tokeye modespec ` | Classic Mirnov mode analysis (vendored [pymodespec](src/tokeye/modespec/classic/PROVENANCE.md), the Python port of the IDL `modespec` tool): power spectrograms, matched-filter toroidal mode-number fits, per-shot mode CSVs. Data fetch needs MDSplus (GA cluster / conda-forge) or a local cache; an example config ships at `src/tokeye/modespec/classic/modes.yaml`. | +| `tokeye elmspec INPUTS...` | ELM detection from the segmentation model's transient channel: per-event time intervals plus per-shot count, ELM frequency (with `--fs`), and duty cycle, written to `elm_events.csv` / `elm_summary.csv`. | +| `tokeye alfvenspec INPUTS...` | Alfvén-eigenmode detection with the `ae_tf_maskrcnn` instance model: per-detection boxes/scores (`ae_detections.csv`) and instance masks. Wide spectrograms are processed in training-width windows automatically. | +| `tokeye eigspec [SCRIPT]` | Interactive modal identification and spectral analysis (vendored [eigspec](src/tokeye/eigspec/PROVENANCE.md), the Python port of the MATLAB toolbox): stochastic subspace ID, AR/PCA, random-projection spectral analysis, clustering (clustering needs `pip install tokeye[eigspec]`). | +| `tokeye modesearch` | Design stage — prints the plan for a searchable database of detected modes. | + +The suite roadmap (including the next-generation `modespec --engine deep`) lives in [docs/ROADMAP.md](docs/ROADMAP.md). + ## Web app guide `tokeye app` (or `python -m tokeye.app`) launches a Gradio interface with three tabs: @@ -152,11 +166,12 @@ This creates a `.venv/`; activate it with `source .venv/bin/activate`, or prefix ## Models -| Registry name | HF file | Description | -| --- | --- | --- | -| `big_tf_unet` | `big_tf_unet_251210.pt` | Transformer U-Net trained on multiscale (multiwindow, multihop) spectrograms. | +| Registry name | HF repo | HF file | Description | +| --- | --- | --- | --- | +| `big_tf_unet` | [`nc1/big_tf_unet`](https://huggingface.co/nc1/big_tf_unet) | `big_tf_unet_251210.pt` | Transformer U-Net trained on multiscale (multiwindow, multihop) spectrograms. | +| `ae_tf_maskrcnn` | `nc1/ae_tf_maskrcnn` | `ae_tf_maskrcnn_251223.pt` | Mask R-CNN instance detector for Alfvén-eigenmode activity (used by `tokeye alfvenspec`). | -Weights are hosted on [Hugging Face](https://huggingface.co/nc1/big_tf_unet) and download automatically the first time a registry name is used (cached in `~/.cache/huggingface`). Override the source repo with the `TOKEYE_HF_REPO` environment variable. +Weights download automatically the first time a registry name is used (cached in `~/.cache/huggingface`). Override the default repo with the `TOKEYE_HF_REPO` environment variable (per-model repos are fixed in the registry). To use a local checkpoint instead, put `.pt`/`.pt2` files in a `model/` directory (picked up by the app's model dropdown) or pass a path directly via `--model PATH`. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md new file mode 100644 index 0000000..bae9b7b --- /dev/null +++ b/docs/ROADMAP.md @@ -0,0 +1,95 @@ +# TokEye roadmap — toward the go-to mode-analysis tool + +TokEye's goal is to cover the jobs DIII-D researchers currently spread across +separate tools (modespec, ad-hoc ELM scripts, per-group AE workflows), so one +install answers "what modes are in this shot?". This file tracks the suite and +collects future ideas worth building. + +## Suite status (0.12.0) + +| Tool | Command | Status | +|------|---------|--------| +| Segmentation | `tokeye run`, `tokeye app` | shipped (big_tf_unet) | +| modespec (classic) | `tokeye modespec ` | shipped — vendored pymodespec (Mirnov n-number fits; needs MDSplus or a cache) | +| modespec (deep) | `tokeye modespec --engine deep` | reserved — single-chord CO2 n-inference, developed in the sibling `integratedmode` project | +| elmspec | `tokeye elmspec INPUTS...` | shipped — ELM events from the transient channel | +| alfvenspec | `tokeye alfvenspec INPUTS...` | shipped, deliberately thin — ae_tf_maskrcnn boxes/masks; awaiting EP-group requirements | +| eigspec | `tokeye eigspec [SCRIPT]` | shipped — vendored (MIT) with import + SSI numeric fixes (see its PROVENANCE.md; fixes worth upstreaming) | +| modesearch | `tokeye modesearch` | design stage — prints the plan | + +## Near-term engineering + +- **Upload `ae_tf_maskrcnn` weights to `nc1/ae_tf_maskrcnn`** (registry entry + and upload-script probe are in place; needs a write-scoped HF token). +- **Upstream the eigspec fixes**: the vendored copy fixes two numeric bugs in + `covariance_driven_ssi` (Hankel channel-interleave, spurious transpose) plus + import-breaking syntax errors — push these back to PlasmaControl/eigspec and + audit the sibling SSI variants (`ssi1ca`, `ssicca`) for the same layout bug. +- **AE weights provenance**: score calibration and a labeled validation set + for alfvenspec before promoting it beyond "runs the model". + +## Mode catalogue schema (the keystone) + +A single record type that every detector emits, so downstream tools compose: + + shot, machine, diagnostic, t_start, t_end, f_low, f_high, + n (nullable), m (nullable), amplitude, confidence, + detector, detector_version, artifact_ref + +- `big_tf_unet` masks → connected regions → records (coherent/transient class) +- `modespec` CSV rows → records with `n` filled +- `elmspec` events → transient records tagged ELM +- `alfvenspec` boxes → records tagged AE + +Once this exists, modesearch is "crawler + storage + filters" rather than a +research project. It also gives papers a uniform unit of comparison across +detectors. + +## modesearch build-out + +1. Crawler: batch job over shot archives (local HDF5 first; MDSplus/toksearch + where reachable) running the suite and emitting catalogue records. +2. Storage: start boring — one parquet/SQLite per campaign; revisit only if + query load demands it. +3. Query CLI: `tokeye modesearch find --n 2 --f 2e3:4e3 --no-elm` → shot list + with matching events. +4. Consumers: the fusion-world-model shot designer learns mode-occurrence + statistics conditioned on plasma parameters; shotsearch intersection + ("shots near this setup that developed a locked mode"). + +## Ideas that would be extremely useful to mode researchers + +- **Mode-number labeling of TokEye masks.** Fuse modespec n-fits with U-Net + regions: overlap a mask region with the (t, f) support of an n-fit and the + region inherits the mode number. Turns "coherent activity" into "n=2 TM", + which is what people actually search for. +- **Mode trajectory tracking.** Follow a detected mode's (f, amplitude, n) + through time: frequency chirps, mode locking (f → 0), rotation braking. + Locked-mode precursors as a first-class query. +- **Cross-diagnostic confirmation.** The same mode seen on Mirnov, CO2, ECE, + and BES with consistent frequency is real; single-diagnostic detections get + a lower confidence. The catalogue schema's `diagnostic` field enables this. +- **ELM database.** elmspec over campaigns → ELM frequency/size statistics vs + pedestal parameters; ELM-free-window finder for AE/TM studies. +- **AE taxonomy.** Classify alfvenspec detections (TAE/RSAE/EAE/BAE) from + frequency-vs-time shape and q-profile context — the EP group's actual need; + gather their requirements before building. +- **Sawtooth/MRE integration.** The vendored classic tree already carries ECE + sawtooth and MRE helpers (`ece_sawteeth.py`, `mre_utils.py`); surface them + as first-class detectors emitting catalogue records. +- **Inter-shot mode.** A between-shots summary (30 s budget): run the suite on + the last shot, print/annotate the mode inventory for the control room. +- **Cross-machine record.** TJ-II validation already exists for the U-Net; + keep the catalogue schema machine-agnostic so C-Mod/NSTX-U/MAST-U archives + can be crawled without schema surgery. +- **Confidence calibration.** Per-detector reliability curves (detected vs + human-labeled) so catalogue confidences are comparable across detectors — + prerequisite for any world-model consumer treating them as probabilities. +- **OMFIT/toksearch hooks.** Thin adapters so existing GA workflows can call + `tokeye.api.TokEye` and the suite CLIs without leaving their environment. + +## Non-goals (for now) + +Real-time control integration (inter-shot is the nearer target), automatic +retraining pipelines, and cross-machine transfer learning beyond what the +existing TJ-II validation demonstrates. diff --git a/pyproject.toml b/pyproject.toml index 2850fcc..a4e5f52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "uv_build" [project] name = "tokeye" -version = "0.11.0" +version = "0.12.0" description = "Automatic classification and localization of fluctuating signals in spectrograms" readme = "README.md" requires-python = ">=3.13" @@ -22,9 +22,15 @@ dependencies = [ "pydantic", "huggingface-hub>=0.30", "tqdm", + "pyyaml", # tokeye.modespec.classic configs (already transitive via gradio) ] [project.optional-dependencies] +# Optional extra for tokeye.eigspec clustering (sklearn imports are +# function-local upstream, so the base package imports without it). +eigspec = [ + "scikit-learn", +] # Mirror of the `train` dependency-group so consumers can `pip install tokeye[train]`. # The training/ablation pipeline modules (src/tokeye/training/) import these; the # core package and the Gradio app do not require them. diff --git a/ruff.toml b/ruff.toml index eb83cec..e8adfc4 100644 --- a/ruff.toml +++ b/ruff.toml @@ -32,4 +32,14 @@ ignore = [ # for research scripts (late imports after sys.path setup, terse one-liners, etc.) "scripts/**" = ["E402", "E702", "E741", "SIM115", "PTH208"] "src/tokeye/extra/eval/**" = ["E702"] +# Vendored code (see PROVENANCE.md in each dir): style rules relaxed to keep +# the upstream diff minimal; correctness rules (F821 etc.) stay active. +"src/tokeye/modespec/classic/**" = [ + "E", "W", "I", "UP", "C4", "FA", "ISC", "ICN", "RET", "SIM", "TID", "TC", + "PTH", "TD", "NPY", "F401", "F841", +] +"src/tokeye/eigspec/**" = [ + "E", "W", "I", "UP", "C4", "FA", "ISC", "ICN", "RET", "SIM", "TID", "TC", + "PTH", "TD", "NPY", "F401", "F841", "F811", "F541", +] diff --git a/scripts/upload_model.py b/scripts/upload_model.py index 6cd48e0..e6767cb 100644 --- a/scripts/upload_model.py +++ b/scripts/upload_model.py @@ -42,15 +42,33 @@ from huggingface_hub import HfApi from huggingface_hub.utils import HfHubHTTPError, LocalTokenNotFoundError -from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID, MODEL_REGISTRY +from tokeye.hub import DEFAULT_MODEL, MODEL_REGISTRY, repo_for if TYPE_CHECKING: from collections.abc import Callable -_PROBE_SHAPE = (1, 1, 64, 64) +def _probe_segmentation(model: nn.Module) -> None: + """Forward pass for the (B, 1, H, W) -> (B, 2, H, W) U-Net contract.""" + model(torch.randn(1, 1, 64, 64)) -def verify_checkpoint(path: Path, builder: Callable[[], nn.Module]) -> None: + +def _probe_rcnn(model: nn.Module) -> None: + """Forward pass for the torchvision R-CNN list-of-images contract.""" + model([torch.randn(3, 64, 64)]) + + +_PROBES: dict[str, Callable[[nn.Module], None]] = { + "big_tf_unet": _probe_segmentation, + "ae_tf_maskrcnn": _probe_rcnn, +} + + +def verify_checkpoint( + path: Path, + builder: Callable[[], nn.Module], + probe: Callable[[nn.Module], None] = _probe_segmentation, +) -> None: """Refuse (via ``SystemExit``) to proceed unless ``path`` is a good checkpoint. "Good" means: a weights-only state dict that loads strictly into a fresh @@ -86,7 +104,7 @@ def verify_checkpoint(path: Path, builder: Callable[[], nn.Module]) -> None: model.eval() with torch.no_grad(): try: - model(torch.randn(*_PROBE_SHAPE)) + probe(model) except Exception as exc: raise SystemExit( f"error: {path} loaded but failed a forward-pass sanity " @@ -123,8 +141,8 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument( "--repo", - default=DEFAULT_REPO_ID, - help="Target Hugging Face Hub repo id (default: %(default)s).", + default=None, + help="Target Hugging Face Hub repo id (default: the model's registry repo).", ) parser.add_argument( "--create", @@ -138,11 +156,13 @@ def main() -> int: args = build_parser().parse_args() spec = MODEL_REGISTRY[args.model] file_path = _resolve_file(args.model, args.file) + if args.repo is None: + args.repo = repo_for(args.model) if not file_path.exists(): raise SystemExit(f"error: checkpoint not found: {file_path}") - verify_checkpoint(file_path, spec.builder) + verify_checkpoint(file_path, spec.builder, _PROBES.get(args.model, _probe_segmentation)) api = HfApi() try: diff --git a/src/tokeye/alfvenspec/__init__.py b/src/tokeye/alfvenspec/__init__.py new file mode 100644 index 0000000..3e77e3f --- /dev/null +++ b/src/tokeye/alfvenspec/__init__.py @@ -0,0 +1,18 @@ +"""Alfvén-eigenmode detection tools (energetic-particle group). + +Deliberately thin for now: ``tokeye alfvenspec`` runs the ``ae_tf_maskrcnn`` +instance-detection model over spectrograms and writes per-detection boxes, +scores, and masks. Deeper EP-group workflows (AE taxonomy, cross-diagnostic +checks) land here once their requirements are gathered — see docs/ROADMAP.md. +""" + +from __future__ import annotations + +from tokeye.alfvenspec.inference import ( + DEFAULT_WINDOW_COLS, + detect, + detect_windowed, + write_detections_csv, +) + +__all__ = ["DEFAULT_WINDOW_COLS", "detect", "detect_windowed", "write_detections_csv"] diff --git a/src/tokeye/alfvenspec/inference.py b/src/tokeye/alfvenspec/inference.py new file mode 100644 index 0000000..d63900a --- /dev/null +++ b/src/tokeye/alfvenspec/inference.py @@ -0,0 +1,131 @@ +"""Inference for the R-CNN detection contract (list of images -> detections). + +Kept separate from :mod:`tokeye.inference`, whose ``model_infer`` assumes the +segmentation contract ``(B, 1, H, W) -> (B, 2, H, W)`` + sigmoid. +""" + +from __future__ import annotations + +import csv +from typing import TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from pathlib import Path + + import torch.nn as nn + + +def detect( + spectrogram: np.ndarray, + model: nn.Module, + *, + score_min: float = 0.5, + mean: float | None = None, + std: float | None = None, +) -> dict[str, np.ndarray]: + """Run an R-CNN detection model on one ``(H, W)`` spectrogram. + + The image is standardized with ``mean``/``std`` (per-sample statistics + when omitted; training used dataset-level stats, so pass them if known) + and fed as a single-channel image — the model's ``GeneralizedRCNNTransform`` + broadcasts it across its 3-element mean/std. Returns numpy arrays: + ``boxes`` (N, 4) xyxy, ``labels`` (N,), ``scores`` (N,), ``masks`` (N, H, W), + filtered to ``scores >= score_min``. + """ + arr = np.asarray(spectrogram, dtype=np.float32) + resolved_mean = float(arr.mean()) if mean is None else mean + resolved_std = float(arr.std()) + 1e-6 if std is None else std + + device = next(model.parameters()).device + image = torch.from_numpy((arr - resolved_mean) / resolved_std) + image = image.unsqueeze(0).float().to(device) + + model.eval() + with torch.no_grad(): + output = model([image])[0] + + keep = output["scores"] >= score_min + return { + "boxes": output["boxes"][keep].cpu().numpy(), + "labels": output["labels"][keep].cpu().numpy(), + "scores": output["scores"][keep].cpu().numpy(), + "masks": output["masks"][keep].squeeze(1).cpu().numpy(), + } + + +DEFAULT_WINDOW_COLS = 710 # training window width; full shots must be windowed +_MIN_WINDOW_COLS = 32 + + +def detect_windowed( + spectrogram: np.ndarray, + model: nn.Module, + *, + window_cols: int = DEFAULT_WINDOW_COLS, + score_min: float = 0.5, + mean: float | None = None, + std: float | None = None, +) -> dict[str, np.ndarray | None]: + """Run :func:`detect` over non-overlapping column windows and merge. + + The R-CNN transform resizes inputs to at most ``max_size`` (1333) columns, + so a full-shot spectrogram (tens of thousands of columns) gets crushed + horizontally and yields nothing; the model was trained on + ~:data:`DEFAULT_WINDOW_COLS`-column views. Box x-coordinates are shifted + back to global columns. ``masks`` is ``None`` whenever more than one + window is used (per-window masks have no common global shape); + ``window_cols <= 0`` disables windowing. + """ + n_cols = spectrogram.shape[1] + if window_cols <= 0 or n_cols <= window_cols: + return detect(spectrogram, model, score_min=score_min, mean=mean, std=std) + + starts = list(range(0, n_cols, window_cols)) + if len(starts) > 1 and n_cols - starts[-1] < _MIN_WINDOW_COLS: + starts.pop() # fold a sliver of a final window into the previous one + + merged: dict[str, list[np.ndarray]] = {"boxes": [], "labels": [], "scores": []} + for index, start in enumerate(starts): + end = starts[index + 1] if index + 1 < len(starts) else n_cols + window = spectrogram[:, start:end] + result = detect(window, model, score_min=score_min, mean=mean, std=std) + result["boxes"][:, [0, 2]] += start + for key in merged: + merged[key].append(result[key]) + + return { + "boxes": np.concatenate(merged["boxes"]), + "labels": np.concatenate(merged["labels"]), + "scores": np.concatenate(merged["scores"]), + "masks": None, + } + + +DETECTION_FIELDS = ("input", "detection", "x1", "y1", "x2", "y2", "label", "score") + + +def write_detections_csv( + path: Path, per_input: list[tuple[str, dict[str, np.ndarray]]] +) -> None: + """One row per detection: box corners (pixel coords), label, score.""" + with path.open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=DETECTION_FIELDS) + writer.writeheader() + for name, detections in per_input: + for index, box in enumerate(detections["boxes"]): + x1, y1, x2, y2 = (float(coord) for coord in box) + writer.writerow( + { + "input": name, + "detection": index, + "x1": x1, + "y1": y1, + "x2": x2, + "y2": y2, + "label": int(detections["labels"][index]), + "score": float(detections["scores"][index]), + } + ) diff --git a/src/tokeye/cli.py b/src/tokeye/cli.py deleted file mode 100644 index 5d9bacd..0000000 --- a/src/tokeye/cli.py +++ /dev/null @@ -1,241 +0,0 @@ -"""``tokeye`` console entry point. - -Argparse only (no new dependencies). Heavy imports (torch, ``tokeye.batch``, -``tokeye.app``) are deferred into each subcommand handler so ``tokeye --help`` -returns instantly and ``tokeye run`` never imports gradio. -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path -from typing import TYPE_CHECKING - -from tokeye.transforms import ( - DEFAULT_CLIP_HIGH, - DEFAULT_CLIP_LOW, - DEFAULT_HOP, - DEFAULT_N_FFT, -) - -if TYPE_CHECKING: - from collections.abc import Sequence - - -def _add_app_subcommand(subparsers: argparse._SubParsersAction) -> None: - parser = subparsers.add_parser("app", help="Launch the TokEye Gradio web app.") - parser.add_argument( - "--port", type=int, default=7860, help="Port to serve the app on." - ) - parser.add_argument( - "--share", action="store_true", help="Create a public Gradio share link." - ) - parser.add_argument( - "--open", - dest="open_browser", - action="store_true", - help="Open the app in a browser on launch.", - ) - parser.set_defaults(handler=_handle_app) - - -def _add_run_subcommand(subparsers: argparse._SubParsersAction) -> None: - parser = subparsers.add_parser("run", help="Run batch inference on one or more inputs.") - parser.add_argument( - "inputs", - nargs="+", - metavar="INPUT", - help="Files, directories of .npy files, or glob patterns.", - ) - parser.add_argument( - "--model", - default=None, - help="Registry name or path to a model checkpoint (default: big_tf_unet).", - ) - parser.add_argument( - "--output-dir", - default="tokeye_output", - help="Directory to write masks and previews to.", - ) - parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT) - parser.add_argument("--hop", type=int, default=DEFAULT_HOP) - parser.add_argument( - "--keep-dc", - action="store_true", - help="Do not clip the DC bin (clipped by default).", - ) - parser.add_argument("--clip-low", type=float, default=DEFAULT_CLIP_LOW) - parser.add_argument("--clip-high", type=float, default=DEFAULT_CLIP_HIGH) - parser.add_argument( - "--log", - action="store_true", - help=( - "Apply log1p to 2D spectrogram inputs stored in linear scale " - "(1D signals are always log-scaled during the STFT)." - ), - ) - parser.add_argument("--threshold", type=float, default=0.5) - parser.add_argument( - "--no-png", - dest="save_png", - action="store_false", - help="Skip PNG overlay previews.", - ) - parser.add_argument("--device", default="auto") - parser.set_defaults(handler=_handle_run) - - -def _add_download_subcommand(subparsers: argparse._SubParsersAction) -> None: - parser = subparsers.add_parser("download", help="Download one or more model checkpoints.") - parser.add_argument( - "models", - nargs="*", - default=None, - metavar="MODEL", - help="Model registry name(s) to download (default: big_tf_unet).", - ) - parser.set_defaults(handler=_handle_download) - - -def _add_example_subcommand(subparsers: argparse._SubParsersAction) -> None: - parser = subparsers.add_parser( - "example", help="Write a synthetic example signal to a .npy file." - ) - parser.add_argument("--output", default="tokeye_example.npy") - parser.add_argument("--duration", type=float, default=2.0) - parser.add_argument("--fs", type=float, default=200_000.0) - parser.add_argument("--seed", type=int, default=0) - parser.set_defaults(handler=_handle_example) - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - prog="tokeye", - description=( - "Automatic classification and localization of fluctuating signals " - "in spectrograms." - ), - ) - parser.add_argument( - "--version", action="store_true", help="Print the tokeye version and exit." - ) - subparsers = parser.add_subparsers(dest="command") - _add_app_subcommand(subparsers) - _add_run_subcommand(subparsers) - _add_download_subcommand(subparsers) - _add_example_subcommand(subparsers) - return parser - - -def _handle_app(args: argparse.Namespace) -> int: - from tokeye.app.__main__ import main as app_main - - app_main(port=args.port, share=args.share, open_browser=args.open_browser) - return 0 - - -def _handle_run(args: argparse.Namespace) -> int: - from huggingface_hub.errors import HfHubHTTPError - - from tokeye import batch - from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID - - stft_kwargs = { - "n_fft": args.n_fft, - "hop": args.hop, - "clip_dc": not args.keep_dc, - "clip_low": args.clip_low, - "clip_high": args.clip_high, - } - model = args.model if args.model is not None else DEFAULT_MODEL - - try: - return batch.run_batch( - args.inputs, - model=model, - out_dir=Path(args.output_dir), - stft_kwargs=stft_kwargs, - save_png=args.save_png, - threshold=args.threshold, - device=args.device, - log=args.log, - ) - except (ValueError, FileNotFoundError) as exc: - hint = ( - " (no data yet? create a demo signal with: tokeye example)" - if "No input files found" in str(exc) - else "" - ) - print(f"error: {exc}{hint}", file=sys.stderr) - return 2 - except (HfHubHTTPError, OSError) as exc: - print( - f"error: could not download model {model!r} from Hugging Face " - f"repo {DEFAULT_REPO_ID!r}: {exc}. If the repo has moved, set " - "TOKEYE_HF_REPO to override.", - file=sys.stderr, - ) - return 2 - - -def _handle_download(args: argparse.Namespace) -> int: - from huggingface_hub.errors import HfHubHTTPError - - from tokeye.hub import DEFAULT_MODEL, DEFAULT_REPO_ID, download_model - - names = args.models or [DEFAULT_MODEL] - for name in names: - try: - path = download_model(name) - except ValueError as exc: - print(f"error: {exc}", file=sys.stderr) - return 2 - except (HfHubHTTPError, OSError) as exc: - print( - f"error: could not download model {name!r} from Hugging Face " - f"repo {DEFAULT_REPO_ID!r}: {exc}. If the repo has moved, set " - "TOKEYE_HF_REPO to override.", - file=sys.stderr, - ) - return 2 - print(path) - return 0 - - -def _handle_example(args: argparse.Namespace) -> int: - import numpy as np - - from tokeye.examples import make_example_signal - - output_path = Path(args.output) - if output_path.suffix != ".npy": - # np.save silently appends ".npy" to paths without that suffix; - # normalize up-front so the printed path is the file that exists. - output_path = output_path.with_suffix(".npy") - output_path.parent.mkdir(parents=True, exist_ok=True) - signal = make_example_signal(duration_s=args.duration, fs=args.fs, seed=args.seed) - np.save(output_path, signal) - print(output_path) - return 0 - - -def main(argv: Sequence[str] | None = None) -> int: - parser = build_parser() - args = parser.parse_args(argv) - - if args.version: - from importlib.metadata import version - - print(version("tokeye")) - return 0 - - if args.command is None: - parser.print_help() - return 2 - - return args.handler(args) - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/tokeye/cli/__init__.py b/src/tokeye/cli/__init__.py new file mode 100644 index 0000000..5d3502d --- /dev/null +++ b/src/tokeye/cli/__init__.py @@ -0,0 +1,75 @@ +"""``tokeye`` console entry point. + +Argparse only (no new dependencies). Heavy imports (torch, ``tokeye.batch``, +``tokeye.app``) are deferred into each subcommand handler so ``tokeye --help`` +returns instantly and ``tokeye run`` never imports gradio. + +Each subcommand lives in its own module under ``tokeye.cli`` and exposes +``add_subcommand(subparsers)``. +""" + +from __future__ import annotations + +import argparse +import sys +from typing import TYPE_CHECKING + +from tokeye.cli import ( + alfvenspec, + app, + download, + eigspec, + elmspec, + example, + modesearch, + modespec, + run, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="tokeye", + description=( + "Automatic classification and localization of fluctuating signals " + "in spectrograms." + ), + ) + parser.add_argument( + "--version", action="store_true", help="Print the tokeye version and exit." + ) + subparsers = parser.add_subparsers(dest="command") + app.add_subcommand(subparsers) + run.add_subcommand(subparsers) + download.add_subcommand(subparsers) + example.add_subcommand(subparsers) + modespec.add_subcommand(subparsers) + elmspec.add_subcommand(subparsers) + alfvenspec.add_subcommand(subparsers) + eigspec.add_subcommand(subparsers) + modesearch.add_subcommand(subparsers) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + + if args.version: + from importlib.metadata import version + + print(version("tokeye")) + return 0 + + if args.command is None: + parser.print_help() + return 2 + + return args.handler(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/tokeye/cli/__main__.py b/src/tokeye/cli/__main__.py new file mode 100644 index 0000000..f40019f --- /dev/null +++ b/src/tokeye/cli/__main__.py @@ -0,0 +1,10 @@ +"""Allow ``python -m tokeye.cli`` to behave like the ``tokeye`` script.""" + +from __future__ import annotations + +import sys + +from tokeye.cli import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/tokeye/cli/_errors.py b/src/tokeye/cli/_errors.py new file mode 100644 index 0000000..e09b7a5 --- /dev/null +++ b/src/tokeye/cli/_errors.py @@ -0,0 +1,17 @@ +"""Shared error reporting for CLI subcommands.""" + +from __future__ import annotations + +import sys + + +def print_hub_error(name: str, exc: Exception) -> None: + """Print a friendly message for a failed Hugging Face model download.""" + from tokeye.hub import repo_for + + print( + f"error: could not download model {name!r} from Hugging Face " + f"repo {repo_for(name)!r}: {exc}. If the repo has moved, set " + "TOKEYE_HF_REPO to override.", + file=sys.stderr, + ) diff --git a/src/tokeye/cli/alfvenspec.py b/src/tokeye/cli/alfvenspec.py new file mode 100644 index 0000000..ae79964 --- /dev/null +++ b/src/tokeye/cli/alfvenspec.py @@ -0,0 +1,164 @@ +"""``tokeye alfvenspec`` — Alfvén-eigenmode detection with ae_tf_maskrcnn.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +from tokeye.transforms import ( + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "alfvenspec", + help="Detect Alfvén-eigenmode activity (boxes + masks via ae_tf_maskrcnn).", + ) + parser.add_argument( + "inputs", + nargs="+", + metavar="INPUT", + help="Files, directories of .npy files, or glob patterns.", + ) + parser.add_argument( + "--model", + default="ae_tf_maskrcnn", + help="Registry name or path to a model checkpoint (default: %(default)s).", + ) + parser.add_argument( + "--output-dir", + default="tokeye_ae", + help="Directory to write detections CSV, masks, and previews to.", + ) + parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT) + parser.add_argument("--hop", type=int, default=DEFAULT_HOP) + parser.add_argument( + "--keep-dc", + action="store_true", + help="Do not clip the DC bin (clipped by default).", + ) + parser.add_argument("--clip-low", type=float, default=DEFAULT_CLIP_LOW) + parser.add_argument("--clip-high", type=float, default=DEFAULT_CLIP_HIGH) + parser.add_argument( + "--log", + action="store_true", + help=( + "Apply log1p to 2D spectrogram inputs stored in linear scale " + "(1D signals are always log-scaled during the STFT)." + ), + ) + parser.add_argument( + "--score-min", + type=float, + default=0.5, + help="Keep detections with at least this score (default: %(default)s).", + ) + parser.add_argument( + "--window-cols", + type=int, + default=710, + help=( + "Process wide spectrograms in windows of this many columns " + "(training width; 0 disables windowing; default: %(default)s)." + ), + ) + parser.add_argument( + "--mean", + type=float, + default=None, + help="Standardization mean (default: per-input statistics).", + ) + parser.add_argument( + "--std", + type=float, + default=None, + help="Standardization std (default: per-input statistics).", + ) + parser.add_argument( + "--no-masks", + dest="save_masks", + action="store_false", + help="Skip writing per-input instance masks (.npy).", + ) + parser.add_argument("--device", default="auto") + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from huggingface_hub.errors import HfHubHTTPError + + from tokeye import batch + from tokeye.alfvenspec import detect_windowed, write_detections_csv + from tokeye.cli._errors import print_hub_error + from tokeye.hub import load_model + + stft_kwargs = { + "n_fft": args.n_fft, + "hop": args.hop, + "clip_dc": not args.keep_dc, + "clip_low": args.clip_low, + "clip_high": args.clip_high, + } + + try: + paths = batch.collect_inputs(args.inputs) + except (ValueError, FileNotFoundError) as exc: + hint = ( + " (no data yet? create a demo signal with: tokeye example)" + if "No input files found" in str(exc) + else "" + ) + print(f"error: {exc}{hint}", file=sys.stderr) + return 2 + + try: + model = load_model(args.model, args.device) + except (ValueError, FileNotFoundError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + except (HfHubHTTPError, OSError) as exc: + print_hub_error(args.model, exc) + return 2 + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + import numpy as np + + all_detections = [] + failures = 0 + for path in paths: + try: + spectrogram = batch.load_input(path, stft_kwargs, log=args.log) + detections = detect_windowed( + spectrogram, + model, + window_cols=args.window_cols, + score_min=args.score_min, + mean=args.mean, + std=args.std, + ) + except Exception as exc: # noqa: BLE001 - mirror `tokeye run`: keep batch going + print(f"error: failed to process {path}: {exc}", file=sys.stderr) + failures += 1 + continue + + all_detections.append((str(path), detections)) + print(f"{path}: {len(detections['boxes'])} detection(s)") + + masks = detections["masks"] # None when the input was windowed + if args.save_masks and masks is not None and len(masks): + np.save(out_dir / f"{path.stem}_ae_masks.npy", masks) + + detections_csv = out_dir / "ae_detections.csv" + write_detections_csv(detections_csv, all_detections) + print(detections_csv) + return failures diff --git a/src/tokeye/cli/app.py b/src/tokeye/cli/app.py new file mode 100644 index 0000000..6574a7b --- /dev/null +++ b/src/tokeye/cli/app.py @@ -0,0 +1,32 @@ +"""``tokeye app`` — launch the Gradio web app.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("app", help="Launch the TokEye Gradio web app.") + parser.add_argument( + "--port", type=int, default=7860, help="Port to serve the app on." + ) + parser.add_argument( + "--share", action="store_true", help="Create a public Gradio share link." + ) + parser.add_argument( + "--open", + dest="open_browser", + action="store_true", + help="Open the app in a browser on launch.", + ) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from tokeye.app.__main__ import main as app_main + + app_main(port=args.port, share=args.share, open_browser=args.open_browser) + return 0 diff --git a/src/tokeye/cli/download.py b/src/tokeye/cli/download.py new file mode 100644 index 0000000..0eb739d --- /dev/null +++ b/src/tokeye/cli/download.py @@ -0,0 +1,41 @@ +"""``tokeye download`` — pre-fetch model checkpoints.""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("download", help="Download one or more model checkpoints.") + parser.add_argument( + "models", + nargs="*", + default=None, + metavar="MODEL", + help="Model registry name(s) to download (default: big_tf_unet).", + ) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from huggingface_hub.errors import HfHubHTTPError + + from tokeye.cli._errors import print_hub_error + from tokeye.hub import DEFAULT_MODEL, download_model + + names = args.models or [DEFAULT_MODEL] + for name in names: + try: + path = download_model(name) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + except (HfHubHTTPError, OSError) as exc: + print_hub_error(name, exc) + return 2 + print(path) + return 0 diff --git a/src/tokeye/cli/eigspec.py b/src/tokeye/cli/eigspec.py new file mode 100644 index 0000000..f23bc7f --- /dev/null +++ b/src/tokeye/cli/eigspec.py @@ -0,0 +1,47 @@ +"""``tokeye eigspec`` — modal identification / spectral analysis (vendored eigspec).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "eigspec", + help=( + "Interactive modal identification and spectral analysis " + "(SSI, AR/PCA, random-projection; MATLAB eigspec port)." + ), + ) + parser.add_argument( + "script", + nargs="?", + default=None, + metavar="SCRIPT", + help="Optional eigspec script file to execute instead of the prompt.", + ) + parser.add_argument( + "--debug", + action="store_true", + help="Enable debug mode with detailed error messages.", + ) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + import os + + import matplotlib as mpl + + mpl.use("Agg") # vendored vis modules import pyplot at module load + + if args.debug: + os.environ["EIGSPEC_DEBUG"] = "1" + + from tokeye.eigspec.cli import EigspecCLI + + EigspecCLI().run(script_file=args.script) + return 0 diff --git a/src/tokeye/cli/elmspec.py b/src/tokeye/cli/elmspec.py new file mode 100644 index 0000000..0266841 --- /dev/null +++ b/src/tokeye/cli/elmspec.py @@ -0,0 +1,185 @@ +"""``tokeye elmspec`` — detect ELM events via the transient-activity channel.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +from tokeye.transforms import ( + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "elmspec", + help="Detect ELM events (transient-channel intervals, count, frequency).", + ) + parser.add_argument( + "inputs", + nargs="+", + metavar="INPUT", + help="Files, directories of .npy files, or glob patterns.", + ) + parser.add_argument( + "--model", + default=None, + help="Registry name or path to a model checkpoint (default: big_tf_unet).", + ) + parser.add_argument( + "--output-dir", + default="tokeye_elms", + help="Directory to write event/summary CSVs (and previews) to.", + ) + parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT) + parser.add_argument("--hop", type=int, default=DEFAULT_HOP) + parser.add_argument( + "--keep-dc", + action="store_true", + help="Do not clip the DC bin (clipped by default).", + ) + parser.add_argument("--clip-low", type=float, default=DEFAULT_CLIP_LOW) + parser.add_argument("--clip-high", type=float, default=DEFAULT_CLIP_HIGH) + parser.add_argument( + "--log", + action="store_true", + help=( + "Apply log1p to 2D spectrogram inputs stored in linear scale " + "(1D signals are always log-scaled during the STFT)." + ), + ) + parser.add_argument( + "--fs", + type=float, + default=None, + help=( + "Sampling rate in Hz of the original signals; enables absolute " + "event times and ELM frequency in the CSVs." + ), + ) + parser.add_argument( + "--threshold", + type=float, + default=0.5, + help="Mask binarization threshold (default: %(default)s).", + ) + parser.add_argument( + "--activity-min", + type=float, + default=0.1, + help=( + "Minimum fraction of active frequency bins for a time column to " + "belong to an ELM (default: %(default)s)." + ), + ) + parser.add_argument( + "--min-gap-cols", + type=int, + default=3, + help="Merge events separated by at most this many columns (default: %(default)s).", + ) + parser.add_argument( + "--min-duration-cols", + type=int, + default=1, + help="Drop events shorter than this many columns (default: %(default)s).", + ) + parser.add_argument( + "--png", + action="store_true", + help="Also write a mask-overlay preview PNG per input.", + ) + parser.add_argument("--device", default="auto") + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from huggingface_hub.errors import HfHubHTTPError + + from tokeye import batch + from tokeye.cli._errors import print_hub_error + from tokeye.elmspec import ( + extract_elm_events, + summarize, + write_events_csv, + write_summary_csv, + ) + from tokeye.hub import DEFAULT_MODEL, load_model + from tokeye.inference import model_infer + + stft_kwargs = { + "n_fft": args.n_fft, + "hop": args.hop, + "clip_dc": not args.keep_dc, + "clip_low": args.clip_low, + "clip_high": args.clip_high, + } + model_name = args.model if args.model is not None else DEFAULT_MODEL + + try: + paths = batch.collect_inputs(args.inputs) + except (ValueError, FileNotFoundError) as exc: + hint = ( + " (no data yet? create a demo signal with: tokeye example)" + if "No input files found" in str(exc) + else "" + ) + print(f"error: {exc}{hint}", file=sys.stderr) + return 2 + + try: + model = load_model(model_name, args.device) + except (ValueError, FileNotFoundError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + except (HfHubHTTPError, OSError) as exc: + print_hub_error(model_name, exc) + return 2 + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + all_events = [] + all_summaries = [] + failures = 0 + for path in paths: + try: + spectrogram = batch.load_input(path, stft_kwargs, log=args.log) + mask = model_infer(spectrogram, model) + events = extract_elm_events( + mask[1], + threshold=args.threshold, + activity_min=args.activity_min, + min_gap_cols=args.min_gap_cols, + min_duration_cols=args.min_duration_cols, + ) + except Exception as exc: # noqa: BLE001 - mirror `tokeye run`: keep batch going + print(f"error: failed to process {path}: {exc}", file=sys.stderr) + failures += 1 + continue + + summary = summarize(events, n_cols=mask.shape[-1], hop=args.hop, fs=args.fs) + all_events.append((str(path), events)) + all_summaries.append((str(path), summary)) + freq = summary["elm_freq_hz"] + freq_text = f", {freq:.1f} Hz" if freq is not None else "" + print(f"{path}: {summary['n_events']} ELM event(s){freq_text}") + + if args.png: + preview_path = out_dir / f"{path.stem}_elm_preview.png" + batch.save_overlay_png(spectrogram, mask, preview_path, threshold=args.threshold) + + events_csv = out_dir / "elm_events.csv" + summary_csv = out_dir / "elm_summary.csv" + write_events_csv(events_csv, all_events, hop=args.hop, fs=args.fs) + write_summary_csv(summary_csv, all_summaries) + print(events_csv) + print(summary_csv) + return failures diff --git a/src/tokeye/cli/example.py b/src/tokeye/cli/example.py new file mode 100644 index 0000000..614eeee --- /dev/null +++ b/src/tokeye/cli/example.py @@ -0,0 +1,37 @@ +"""``tokeye example`` — write a synthetic demo signal.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "example", help="Write a synthetic example signal to a .npy file." + ) + parser.add_argument("--output", default="tokeye_example.npy") + parser.add_argument("--duration", type=float, default=2.0) + parser.add_argument("--fs", type=float, default=200_000.0) + parser.add_argument("--seed", type=int, default=0) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + import numpy as np + + from tokeye.examples import make_example_signal + + output_path = Path(args.output) + if output_path.suffix != ".npy": + # np.save silently appends ".npy" to paths without that suffix; + # normalize up-front so the printed path is the file that exists. + output_path = output_path.with_suffix(".npy") + output_path.parent.mkdir(parents=True, exist_ok=True) + signal = make_example_signal(duration_s=args.duration, fs=args.fs, seed=args.seed) + np.save(output_path, signal) + print(output_path) + return 0 diff --git a/src/tokeye/cli/modesearch.py b/src/tokeye/cli/modesearch.py new file mode 100644 index 0000000..ae83209 --- /dev/null +++ b/src/tokeye/cli/modesearch.py @@ -0,0 +1,37 @@ +"""``tokeye modesearch`` — mode database (design stage, prints the vision).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + +_DESCRIPTION = """\ +modesearch is not implemented yet. The plan: + + 1. An offline crawler runs the TokEye suite (big_tf_unet, modespec, + elmspec, alfvenspec) over shot archives and indexes every detected + mode: shot, time interval, frequency band, mode numbers, amplitude, + detector provenance. + 2. Researchers query the index -- e.g. "shots with an n=2 tearing mode + between 2-4 kHz during an ELM-free period" -- instead of re-scanning + raw data. + 3. The same index feeds the fusion-world-model shot designer with mode + occurrence statistics. + +Design notes: src/tokeye/modesearch/README.md and docs/ROADMAP.md. +""" + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "modesearch", + help="Mode database + queries (design stage; prints the plan).", + ) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + print(_DESCRIPTION) + return 0 diff --git a/src/tokeye/cli/modespec.py b/src/tokeye/cli/modespec.py new file mode 100644 index 0000000..08788a3 --- /dev/null +++ b/src/tokeye/cli/modespec.py @@ -0,0 +1,57 @@ +"""``tokeye modespec`` — classic Mirnov mode-number analysis (vendored pymodespec).""" + +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser( + "modespec", + help=( + "Classic DIII-D mode analysis: Mirnov spectrograms, toroidal " + "mode-number fits, per-shot mode CSVs (needs MDSplus or a cache)." + ), + ) + parser.add_argument( + "config", + metavar="CONFIG", + help="YAML config listing shots and analysis parameters (see modes.yaml).", + ) + parser.add_argument( + "--engine", + choices=["classic"], + default="classic", + help="Analysis engine (only 'classic' today; 'deep' is planned).", + ) + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from pathlib import Path + + import matplotlib as mpl + + mpl.use("Agg") # vendored modules import pyplot at module load + + from tokeye.modespec.classic import run_config + + config_path = Path(args.config) + if not config_path.exists(): + example = Path(__file__).parent.parent / "modespec" / "classic" / "modes.yaml" + print( + f"error: config not found: {config_path} " + f"(example config: {example})", + file=sys.stderr, + ) + return 2 + + try: + return run_config(config_path) + except (KeyError, ValueError) as exc: + print(f"error: bad config {config_path}: {exc}", file=sys.stderr) + return 2 diff --git a/src/tokeye/cli/run.py b/src/tokeye/cli/run.py new file mode 100644 index 0000000..419c241 --- /dev/null +++ b/src/tokeye/cli/run.py @@ -0,0 +1,103 @@ +"""``tokeye run`` — headless batch inference.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +from tokeye.transforms import ( + DEFAULT_CLIP_HIGH, + DEFAULT_CLIP_LOW, + DEFAULT_HOP, + DEFAULT_N_FFT, +) + +if TYPE_CHECKING: + import argparse + + +def add_subcommand(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("run", help="Run batch inference on one or more inputs.") + parser.add_argument( + "inputs", + nargs="+", + metavar="INPUT", + help="Files, directories of .npy files, or glob patterns.", + ) + parser.add_argument( + "--model", + default=None, + help="Registry name or path to a model checkpoint (default: big_tf_unet).", + ) + parser.add_argument( + "--output-dir", + default="tokeye_output", + help="Directory to write masks and previews to.", + ) + parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT) + parser.add_argument("--hop", type=int, default=DEFAULT_HOP) + parser.add_argument( + "--keep-dc", + action="store_true", + help="Do not clip the DC bin (clipped by default).", + ) + parser.add_argument("--clip-low", type=float, default=DEFAULT_CLIP_LOW) + parser.add_argument("--clip-high", type=float, default=DEFAULT_CLIP_HIGH) + parser.add_argument( + "--log", + action="store_true", + help=( + "Apply log1p to 2D spectrogram inputs stored in linear scale " + "(1D signals are always log-scaled during the STFT)." + ), + ) + parser.add_argument("--threshold", type=float, default=0.5) + parser.add_argument( + "--no-png", + dest="save_png", + action="store_false", + help="Skip PNG overlay previews.", + ) + parser.add_argument("--device", default="auto") + parser.set_defaults(handler=_handle) + + +def _handle(args: argparse.Namespace) -> int: + from huggingface_hub.errors import HfHubHTTPError + + from tokeye import batch + from tokeye.cli._errors import print_hub_error + from tokeye.hub import DEFAULT_MODEL + + stft_kwargs = { + "n_fft": args.n_fft, + "hop": args.hop, + "clip_dc": not args.keep_dc, + "clip_low": args.clip_low, + "clip_high": args.clip_high, + } + model = args.model if args.model is not None else DEFAULT_MODEL + + try: + return batch.run_batch( + args.inputs, + model=model, + out_dir=Path(args.output_dir), + stft_kwargs=stft_kwargs, + save_png=args.save_png, + threshold=args.threshold, + device=args.device, + log=args.log, + ) + except (ValueError, FileNotFoundError) as exc: + hint = ( + " (no data yet? create a demo signal with: tokeye example)" + if "No input files found" in str(exc) + else "" + ) + print(f"error: {exc}{hint}", file=sys.stderr) + return 2 + except (HfHubHTTPError, OSError) as exc: + print_hub_error(model, exc) + return 2 diff --git a/src/tokeye/eigspec/LICENSE.md b/src/tokeye/eigspec/LICENSE.md new file mode 100644 index 0000000..c4954f1 --- /dev/null +++ b/src/tokeye/eigspec/LICENSE.md @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 PlasmaControl + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/tokeye/eigspec/PROVENANCE.md b/src/tokeye/eigspec/PROVENANCE.md new file mode 100644 index 0000000..96722a6 --- /dev/null +++ b/src/tokeye/eigspec/PROVENANCE.md @@ -0,0 +1,47 @@ +# Vendored code provenance + +- Upstream: `git@github.com:PlasmaControl/eigspec.git` (public) +- Pinned commit: `923ad5da97ed8c69cf83c99e0163510c6abe20c2` ("Add MIT License to the project") +- Vendored: 2026-07-06 +- License: MIT (see `LICENSE.md`, copied from upstream) + +## Files taken + +Everything under upstream `src/eigspec/` (the package code: `analysis/`, +`io/`, `utils/`, `vis/`, `cli.py`, `__init__.py`) plus `LICENSE.md`. +Upstream `matlab/`, `assets/`, `demo/`, `examples/`, `docs/`, tests, and +figure PNGs were not vendored. + +## Local modifications + +Upstream could not be imported at all — `lambda` used as an attribute name is +a hard SyntaxError, and several annotations referenced names that were never +imported. Fixes: + +- `utils/data_extraction.py` (2 sites) and `vis/spectral_plots.py` (5 sites): + `block.mrep.m0.lambda` → `getattr(..., 'lambda')`. +- `utils/subspace_identification.py`: added the missing + `Any, Dict, Tuple, Union` typing imports and `import numpy.typing as npt` + (annotations referenced them → `NameError` at import on Python 3.13). +- `utils/subspace_identification.py` — two numeric bugs in + `covariance_driven_ssi` (worth upstreaming): + 1. The block-Hankel matrix was flattened channel-major + (`data_block.T.flatten()`), but every downstream slice + (`[:m*p]`, `[m:m*f]`, `[:m*(f-1)]`) assumes time-block-major rows like + the MATLAB original — past/future blocks were scrambled and recovered + pole frequencies were wrong (e.g. 2x for a 2-channel sin/cos pair). + Fixed by flattening the (time, channel) block in C order. + 2. The `A = O1 \\ O2` step carried a spurious `.T` ("transpose to match + MATLAB"): `np.linalg.lstsq(O1, O2)[0]` is already `pinv(O1) @ O2`, + identical to backslash. Eigenvalues survive a transpose but mode + shapes do not. Removed (both single- and multi-order paths). + After the fixes, a damped-sinusoid test recovers both the pole frequency + and the damping ratio to <1% (see tests/test_eigspec.py). The sibling + `canonical_correlation_ssi` / `ssi1ca` / `ssicca` functions were NOT + audited for the same Hankel-layout issue. +- `PROVENANCE.md` (this file) is an addition, not an upstream file. +- Style rules are relaxed for this directory in `ruff.toml` + (vendored-code policy); correctness rules (F821 etc.) remain active. + +Re-vendoring: clone upstream at a newer commit, re-copy `src/eigspec/*`, +re-apply (or upstream) the fixes above, and update the pinned commit. diff --git a/src/tokeye/eigspec/__init__.py b/src/tokeye/eigspec/__init__.py new file mode 100644 index 0000000..eb8ff56 --- /dev/null +++ b/src/tokeye/eigspec/__init__.py @@ -0,0 +1,229 @@ +""" +eigspec - Python port of eigspec MATLAB toolbox for spectral analysis and modal identification + +This package provides tools for: +- Stochastic Subspace Identification (SSI) +- AR/PCA time-series modeling +- Random projection spectral analysis +- Modal analysis and shape estimation +- Block-based processing for large datasets +- Clustering analysis for modal pattern recognition +- Comprehensive visualization and data I/O + +Python port of the MATLAB eigspec toolbox developed for plasma physics applications. +Main MATLAB entry points correspond to: +- eigspec_mmain.m - Main analysis workflow +- rndspecx.m - Random projection spectral analysis +- view_pcaspec_*.m - Visualization and results analysis +- clus_*.m - Clustering analysis functions +- ssi*.m - Subspace identification algorithms +""" + +__version__ = "0.1.0" + +# Main analysis functions +from .analysis.random_projection import RandomProjectionSpectralAnalysisResult, random_projection_spectral_analysis + +# Core modal analysis utilities +from .utils.modal_analysis import ModalList, ModalShortlist, ShapeEstimates, extract_modal_parameters, order_mac, shapes_from_freq, complex_vector_scalar_fit, shape2mn + +# Block processing functions +from .utils.block_processing import ( + BlockAnalysisResult, + RandomProjectionResult, + random_projection_block_analysis +) + +# System identification algorithms +from .utils.subspace_identification import ( + covariance_driven_ssi, + canonical_correlation_ssi, + ssi1ca, + ssicca, + SubspaceIdentificationResult, + StateSpaceModel +) + +# AR/PCA modeling +from .utils.autoregressive_pca import arpca, ARPCAResult, ARPCAModel + +# Signal processing utilities +from .utils.signal_processing import ( + FFTSpectralResult, + CoherenceResult, + AssessmentResult, + CorrelationAssessmentResult, + compute_zero_mean_spectrum, + aggregate_fft, + filter_signal, + create_window, + fftspec, + fftspec1, + fftspecwin, + zmfftspec, + yintegrate, + ydecimate, + yresample, + yaddgauss, + coherence_filter, + ar_assessment, + correlation_assessment, + arassess, + corrassess, + fdmspec1, + kdftspec, + yfilt, + yinterpolate, + qplot_data +) + +# MATLAB utility functions +from .utils.matlab_utilities import ( + kfoldcov, + logdet, + srteig, + zpdftmatrix, + fdm1dk, + weighted_rms +) + +# FDM analysis functions +from .utils.fdm_analysis import ( + fdm1d, + rndspec, + pcaspecx +) + +# Data extraction utilities +from .utils.data_extraction import ( + extract_ptrefs, + collect_rep_data, + PointReference, + PointReferences +) + +# Clustering analysis utilities +from .utils.clustering import ( + DistanceMetric, + ClusteringResult, + distance_matrix, + similarity_matrix, + kmeans_clustering, + spectral_clustering, + medoid_clustering, + mac_value, + trim_cluster_mac, + clus_similarity_matrix, + clus_distance_matrix, + spclus_spectral, + spclus_knn_similarity_matrix, + clus_krnn_enhance, +) + +# Import visualization module if available +try: + from . import vis + _VIS_AVAILABLE = True +except ImportError: + _VIS_AVAILABLE = False + +# Import I/O module if available +try: + from . import io + _IO_AVAILABLE = True +except ImportError: + _IO_AVAILABLE = False + +__all__ = [ + # Main analysis entry points + 'random_projection_spectral_analysis', + 'RandomProjectionSpectralAnalysisResult', + + # Modal analysis data structures + 'ModalList', + 'ModalShortlist', + 'ShapeEstimates', + + # Block processing + 'BlockAnalysisResult', + 'RandomProjectionResult', + 'random_projection_block_analysis', + + # Core analysis functions + 'extract_modal_parameters', +'order_mac', +'shapes_from_freq', +'complex_vector_scalar_fit', +'shape2mn', +'modal_mac_matrix', +'sort_modes_by_frequency', +'normalize_mode_shapes', +'mode_shape_scaling_factor', +'modal_correlation_coefficient', +'extract_modal_parameters', + + # System identification + 'covariance_driven_ssi', + 'canonical_correlation_ssi', + 'ssi1ca', + 'ssicca', + 'SubspaceIdentificationResult', + 'StateSpaceModel', + + # AR/PCA modeling + 'arpca', + 'ARPCAResult', + 'ARPCAModel', + + # Signal processing + 'FFTSpectralResult', + 'CoherenceResult', + 'AssessmentResult', + 'CorrelationAssessmentResult', + 'compute_zero_mean_spectrum', + 'aggregate_fft', + 'filter_signal', + 'create_window', + 'fftspec', + 'coherence_filter', + 'ar_assessment', + 'correlation_assessment', + + # MATLAB utilities + 'kfoldcov', + 'logdet', + 'srteig', + 'zpdftmatrix', + 'fdm1dk', + 'weighted_rms', + + # FDM analysis + 'fdm1d', + 'rndspec', + 'pcaspecx', + + # Data extraction + 'extract_ptrefs', + 'collect_rep_data', + 'PointReference', + 'PointReferences', + + # Clustering analysis + 'DistanceMetric', + 'ClusteringResult', + 'distance_matrix', + 'similarity_matrix', + 'kmeans_clustering', + 'spectral_clustering', + 'medoid_clustering', + 'mac_value', + 'trim_cluster_mac', +] + +# Add visualization to exports if available +if _VIS_AVAILABLE: + __all__.append('vis') + +# Add I/O to exports if available +if _IO_AVAILABLE: + __all__.append('io') \ No newline at end of file diff --git a/src/tokeye/eigspec/analysis/__init__.py b/src/tokeye/eigspec/analysis/__init__.py new file mode 100644 index 0000000..1539cb7 --- /dev/null +++ b/src/tokeye/eigspec/analysis/__init__.py @@ -0,0 +1,3 @@ +""" +Main scripts for eigspec package. +""" \ No newline at end of file diff --git a/src/tokeye/eigspec/analysis/random_projection.py b/src/tokeye/eigspec/analysis/random_projection.py new file mode 100644 index 0000000..a91ff70 --- /dev/null +++ b/src/tokeye/eigspec/analysis/random_projection.py @@ -0,0 +1,194 @@ +""" +Random projection spectral analysis for eigspec package. + +This module provides the main random projection spectral analysis functionality: +- High-level interface for multi-channel time series analysis +- Automated modal identification with random dimensionality reduction +- Comprehensive analysis workflow combining SSI, AR-PCA, and modal analysis + +Based on the MATLAB eigspec toolbox main analysis functions: +- rndspecx.m - Core random projection spectral analysis algorithm +- eigspec_mmain.m - Main analysis entry point and workflow +- view_pcaspec_results.m - Results processing and classification +- view_pcaspec_prototypes.m - Prototype-based modal analysis +""" + +import time +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import numpy as np +from numpy.typing import NDArray + +from ..utils.block_processing import random_projection_block_analysis +from ..utils.utils import demean + +@dataclass +class RandomProjectionBlockResult: + """ + Container for individual block processing results in spectral analysis. + + Attributes: + reduced_dimension_matrix: Modal analysis results + reduced_dimension_array: Shape estimates + projection_matrix: Random projection matrix used for this block + processing_time: Time taken to process block in seconds + demean_block: Whether block was demeaned before processing + time_step: Time step between samples in seconds + time_slice: (start_idx, end_idx) time indices of block + centre_time: Center time of block in seconds + filter_time: Filter time in seconds + """ + reduced_dimension_matrix: object # ModalShortlist + reduced_dimension_array: object # ShapeEstimates or None + projection_matrix: Optional[NDArray[np.floating]] # Random projection matrix + processing_time: float + demean_block: bool + time_step: float + time_slice: Tuple[int, int] + centre_time: float + filter_time: float + + # Add property aliases for backward compatibility with tests + @property + def modal_report(self) -> object: + """Alias for reduced_dimension_matrix for test compatibility.""" + return self.reduced_dimension_matrix + +@dataclass +class RandomProjectionSpectralAnalysisResult: + """ + Container for spectral analysis results. + + Attributes: + block_results: List of results for each processed block + total_processing_time: Total processing time in seconds + reduced_dimension: Analysis parameters: + For SSI: [reduced_dim, future, past, order1, order2] + For AR/PCA: [reduced_dim, past, order1, order2] + where: + - reduced_dim: Target dimension (<0 for orthonormal, >0 for random, 0 for no projection) + - future/past: Number of future/past samples for SSI + - order1/order2: Model orders + block_parameters: (block_size, block_stride) for block processing + threshold_parameters: (MAC threshold, DST threshold) for mode matching + use_canonical_correlation_analysis: Whether CCA/CVA was used in SSI + analysis_name: Name of analysis routine used + """ + block_results: List[RandomProjectionBlockResult] + total_processing_time: float + reduced_dimension: List[int] + block_parameters: Tuple[int, int] + threshold_parameters: Tuple[float, float] + use_canonical_correlation_analysis: bool + analysis_name: str + +def random_projection_spectral_analysis( + time_array: Union[None, NDArray[np.number], float, int], + signal_array: NDArray[np.number], + block_parameters: Tuple[int, int], + reduced_dimension: List[int], + threshold_parameters: Tuple[float, float], + use_canonical_correlation_analysis: bool = False, + random_seed: Optional[int] = None +) -> RandomProjectionSpectralAnalysisResult: + """ + Block-based spectral analysis of multivariate time-series using random projections. + + Args: + time_array: Time vector, scalar 0, or None. If None or 0, uses sample indices. + signal_array: Input data matrix, shape (n_samples, n_channels) + block_parameters: (block_size, block_stride) for block processing + reduced_dimension: Analysis parameters: + For SSI: [reduced_dim, future, past, order1, order2] + For AR/PCA: [reduced_dim, past, order1, order2] + where: + - reduced_dim: Target dimension (<0 for orthonormal, >0 for random, 0 for no projection) + - future/past: Number of future/past samples for SSI + - order1/order2: Model orders + threshold_parameters: (MAC threshold, DST threshold) for mode matching + use_canonical_correlation_analysis: Whether to use CCA/CVA in SSI + random_seed: Optional seed for random number generation + + Returns: + RandomProjectionSpectralAnalysisResult containing analysis results + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*t)[:, np.newaxis] + >>> block_params = (100, 50) # 100-sample blocks with 50-sample stride + >>> reduced_dim = [-5, 10, 10, 2, 2] # SSI with orthonormal projection to 5D + >>> thresh = (0.9, 0.9) # MAC and distance thresholds + >>> result = random_projection_spectral_analysis(t, y, block_params, reduced_dim, thresh) + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + if not isinstance(reduced_dimension, list) or not all(isinstance(x, int) for x in reduced_dimension): + raise TypeError("reduced_dimension must be a list of integers") + if len(reduced_dimension) not in (4, 5): + raise ValueError("reduced_dimension must have 4 elements (AR/PCA) or 5 elements (SSI)") + + # Set random seed if provided + if random_seed is not None: + np.random.seed(random_seed) + + n, _ = signal_array.shape + if time_array is None or (isinstance(time_array, (int, float)) and time_array == 0): + time_vector = np.arange(n) + time_step = -1 + else: + time_vector = np.asarray(time_array).reshape(-1) + if len(time_vector) != n: + raise ValueError("Length of time_array does not match length of signal_array") + time_step = time_vector[1] - time_vector[0] + + block_size, block_stride = block_parameters + block_start_indices = np.arange(0, n + 1 - block_size, block_stride) + n_block = len(block_start_indices) + + # Initialize results + block_results: List[RandomProjectionBlockResult] = [] + + ttl = time.time() + + for block_index in range(n_block): + time_start = block_start_indices[block_index] + time_end = time_start + block_size - 1 + + # Process block, time-slice [time_start, time_end] + start_time = time.time() + + # Demean and process block + demeaned_block = demean(signal_array[time_start:time_end+1, :]) + random_projection_result = random_projection_block_analysis( + demeaned_block, reduced_dimension, threshold_parameters, + use_canonical_correlation_analysis, random_seed + ) + + # Store results + # Process all the channels in one go - if can parallelize, then don't need to do random projection + block_result = RandomProjectionBlockResult( + reduced_dimension_matrix=random_projection_result.modal_analysis, + reduced_dimension_array=random_projection_result.shape_estimates, + projection_matrix=random_projection_result.projection_matrix, + processing_time=time.time() - start_time, + demean_block=True, + time_step=time_step, + time_slice=(time_start, time_end), + centre_time=float((time_vector[time_start] + time_vector[time_end]) / 2), + filter_time=float(time_vector[time_end]) + ) + block_results.append(block_result) + + # Create final results + result = RandomProjectionSpectralAnalysisResult( + block_results=block_results, + total_processing_time=time.time() - ttl, + reduced_dimension=reduced_dimension, + block_parameters=block_parameters, + threshold_parameters=threshold_parameters, + use_canonical_correlation_analysis=use_canonical_correlation_analysis, + analysis_name=f"random_projection_spectral_analysis({'ssi' if len(reduced_dimension) == 5 else 'ar/pca'})" + ) + + return result diff --git a/src/tokeye/eigspec/cli.py b/src/tokeye/eigspec/cli.py new file mode 100644 index 0000000..8adac5d --- /dev/null +++ b/src/tokeye/eigspec/cli.py @@ -0,0 +1,675 @@ +#!/usr/bin/env python3 +""" +Command Line Interface for eigspec - Python port of eigspec_mmain.m + +This module provides an interactive command-line interface for spectral analysis +and modal identification, replicating the functionality of the MATLAB eigspec_mmain.m. + +Features: +- Interactive command prompt with help system +- Script execution capability +- Data loading and preprocessing +- Spectral analysis workflows +- Assessment and clustering tools +- Visualization and export functions + +Commands mirror the MATLAB version where possible for familiarity. +""" + +import argparse +import os +import sys +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Any, Union + +import numpy as np + +# eigspec imports +from . import __version__ +from .utils.signal_processing import ar_assessment, correlation_assessment, get_performance_info +from .utils.modal_analysis import complex_vector_scalar_fit, modal_mac_matrix +from .analysis.random_projection import random_projection_spectral_analysis + + +class EigspecCLI: + """Main command-line interface for eigspec analysis.""" + + def __init__(self) -> None: + """Initialize the CLI with default settings.""" + self.prompt = "-->" + self.data: Dict[str, Any] = {} + self.settings: Dict[str, Any] = { + # Default analysis parameters (equivalent to MATLAB defaults) + 'fft_options': { + 'bss': [512, 128], # [block_size, block_stride] + 'nfft': 2048, + 'nsmooth': 5, + 'rpdim': 0, + 'winstr': 'hamming' + }, + 'rndspec_options': { + 'bss': [500, 400], + 'rfpn': [-12, 10, 20, 12, 20], # random projection parameters + 'thresh': [0.998, 0.998] + }, + 'clus_options': { + 'numrestarts': 10, + 'kr': [40, 3], # k-NN parameters + 'knn': 40 + }, + 'MN_MMAX': 8, # Maximum toroidal mode number + 'MN_NMAX': 5, # Maximum poloidal mode number + 'default_ar_lag': 12, + 'default_filter_order': 2 + } + + # Analysis results storage + self.current_data = None # Raw data + self.processed_data = None # Preprocessed data + self.analysis_results = None # Spectral analysis results + self.assessment_results = None # Channel assessment results + self.cluster_results = None # Clustering results + + # Command registry + self.commands: Dict[str, callable] = { + 'help': self.cmd_help, + 'quit': self.cmd_quit, + 'exit': self.cmd_quit, + 'version': self.cmd_version, + 'status': self.cmd_status, + 'settings': self.cmd_settings, + 'load': self.cmd_load, + 'assess': self.cmd_assess, + 'corr': self.cmd_corr, + 'rndspec': self.cmd_rndspec, + 'view': self.cmd_view, + 'export': self.cmd_export, + 'script': self.cmd_script, + 'clear': self.cmd_clear + } + + self.running = True + self.script_mode = False + self.script_commands: List[str] = [] + self.script_index = 0 + + def run(self, script_file: Optional[str] = None) -> None: + """Run the CLI interface. + + Args: + script_file: Optional script file to execute on startup + """ + self._print_banner() + + if script_file: + self.cmd_script([script_file]) + + while self.running: + try: + if self.script_mode: + if self.script_index < len(self.script_commands): + command_line = self.script_commands[self.script_index] + print(f"[@script/line #{self.script_index+1}]{self.prompt}{command_line}") + self.script_index += 1 + else: + print("[script done.]") + self.script_mode = False + continue + else: + command_line = input(f"{self.prompt}").strip() + + if command_line: + self._execute_command(command_line) + + except KeyboardInterrupt: + print("\nInterrupted. Type 'quit' to exit.") + continue + except EOFError: + print("\nExiting...") + break + except Exception as e: + print(f"Error: {e}") + if os.getenv('EIGSPEC_DEBUG'): + traceback.print_exc() + + def _print_banner(self) -> None: + """Print startup banner.""" + import platform + + print("<<< entering eigspec: (array magnetics fluctuation analysis)") + print(f"<<< version: {__version__} (Python)") + print(f"<<< platform: {platform.system()} {platform.release()}") + print(f"<<< working directory: {os.getcwd()}") + print("<<< type 'help' for information, or 'quit' to exit") + + # Show performance info + perf_info = get_performance_info() + if perf_info['numba_available']: + print("<<< performance: Numba JIT compilation available") + if perf_info['joblib_available']: + print("<<< performance: parallel processing available") + + def _execute_command(self, command_line: str) -> None: + """Execute a command line.""" + # Skip comments + if command_line.startswith('//'): + return + + # Parse command and arguments + parts = command_line.split() + if not parts: + return + + command = parts[0].lower() + args = parts[1:] if len(parts) > 1 else [] + + if command in self.commands: + try: + self.commands[command](args) + except Exception as e: + print(f"Error executing '{command}': {e}") + if os.getenv('EIGSPEC_DEBUG'): + traceback.print_exc() + else: + print(f"Unknown command: {command}. Type 'help' for available commands.") + + def cmd_help(self, args: List[str]) -> None: + """Show help information.""" + if not args: + print("eigspec - Spectral analysis and modal identification") + print("\nAvailable commands:") + print(" Data I/O:") + print(" load - Load data from file") + print(" export - Export results") + print(" clear - Clear loaded data") + print("") + print(" Analysis:") + print(" assess [bsize] [bstride] - AR-based channel assessment") + print(" corr [bsize] [bstride] - Correlation assessment") + print(" rndspec [options] - Random projection spectral analysis") + print("") + print(" Utilities:") + print(" view [results] - View analysis results") + print(" settings [param] [value] - Show/set analysis parameters") + print(" status - Show current status") + print(" script - Execute script file") + print("") + print(" System:") + print(" version - Show version information") + print(" help [command] - Show help") + print(" quit - Exit program") + else: + command = args[0].lower() + if command == 'load': + print("load - Load data from file") + print(" Supported formats: .npy, .npz, .csv, .txt") + print(" Data should be organized as (n_samples, n_channels)") + elif command == 'assess': + print("assess [block_size] [block_stride] - AR-based assessment") + print(" ar_lag: AR model lag order (e.g., 12)") + print(" block_size: Analysis block size (default: 800)") + print(" block_stride: Block stride (default: 800)") + elif command == 'rndspec': + print("rndspec [reduced_dim] [future] [past] [order1] [order2] - Spectral analysis") + print(" All parameters optional, uses settings defaults if not provided") + print(" reduced_dim: Random projection dimension (<0=orthonormal, >0=random)") + print(" future/past: SSI horizon parameters") + print(" order1/order2: Model orders") + else: + print(f"No detailed help available for '{command}'") + + def cmd_quit(self, args: List[str]) -> None: + """Exit the program.""" + print("Goodbye!") + self.running = False + + def cmd_version(self, args: List[str]) -> None: + """Show version information.""" + import platform + + print(f"eigspec version: {__version__}") + print(f"Python version: {platform.python_version()}") + print(f"Platform: {platform.system()} {platform.release()}") + + perf_info = get_performance_info() + print(f"Numba available: {perf_info['numba_available']}") + print(f"Joblib available: {perf_info['joblib_available']}") + + def cmd_status(self, args: List[str]) -> None: + """Show current analysis status.""" + print("=== Analysis Status ===") + + if self.current_data is not None: + n_samples, n_channels = self.current_data.shape + print(f"Data loaded: {n_samples} samples, {n_channels} channels") + else: + print("Data: None loaded") + + if self.analysis_results is not None: + n_blocks = len(self.analysis_results.block_results) + print(f"Spectral analysis: {n_blocks} blocks processed") + else: + print("Spectral analysis: Not performed") + + if self.assessment_results is not None: + print("Assessment: Available") + else: + print("Assessment: Not performed") + + def cmd_settings(self, args: List[str]) -> None: + """Show or modify analysis settings.""" + if not args: + print("=== Current Settings ===") + for category, params in self.settings.items(): + print(f"{category}:") + if isinstance(params, dict): + for key, value in params.items(): + print(f" {key}: {value}") + else: + print(f" {params}") + elif len(args) == 1: + # Show specific setting + param = args[0] + found = False + for category, params in self.settings.items(): + if isinstance(params, dict) and param in params: + print(f"{category}.{param}: {params[param]}") + found = True + elif param == category: + print(f"{category}: {params}") + found = True + if not found: + print(f"Setting '{param}' not found") + else: + print("Usage: settings [parameter] [value]") + + def cmd_load(self, args: List[str]) -> None: + """Load data from file.""" + if not args: + print("Usage: load ") + return + + filename = args[0] + + try: + if filename.endswith('.npy'): + data = np.load(filename) + elif filename.endswith('.npz'): + archive = np.load(filename) + # Try common key names + for key in ['data', 'signals', 'y', 'Y']: + if key in archive: + data = archive[key] + break + else: + # Use first array found + data = next(iter(archive.values())) + elif filename.endswith(('.csv', '.txt')): + data = np.loadtxt(filename, delimiter=',') + else: + print(f"Unsupported file format: {filename}") + return + + # Ensure 2D array + if data.ndim == 1: + data = data.reshape(-1, 1) + + self.current_data = data + n_samples, n_channels = data.shape + print(f"Loaded {filename}: {n_samples} samples, {n_channels} channels") + + # Clear previous analysis results + self.analysis_results = None + self.assessment_results = None + self.cluster_results = None + + except Exception as e: + print(f"Failed to load {filename}: {e}") + + def cmd_assess(self, args: List[str]) -> None: + """Perform AR-based channel assessment.""" + if self.current_data is None: + print("No data loaded. Use 'load ' first.") + return + + if not args: + # Use default parameters + ar_lag = self.settings['default_ar_lag'] + block_size = self.settings['rndspec_options']['bss'][0] + block_stride = self.settings['rndspec_options']['bss'][1] + else: + try: + ar_lag = int(args[0]) + block_size = int(args[1]) if len(args) > 1 else self.settings['rndspec_options']['bss'][0] + block_stride = int(args[2]) if len(args) > 2 else self.settings['rndspec_options']['bss'][1] + except ValueError: + print("Usage: assess [block_size] [block_stride]") + return + + print(f"AR-based assessment: lag={ar_lag}, block=[{block_size}, {block_stride}]") + + try: + self.assessment_results = ar_assessment( + time_vector=None, + signal_array=self.current_data, + ar_lag=ar_lag, + block_parameters=(block_size, block_stride) + ) + + print("Assessment complete. Results:") + print(f" Median predictability: {self.assessment_results.median_scores[:, 0].mean():.3f}") + print(f" Median participation: {self.assessment_results.median_scores[:, 1].mean():.3f}") + print("Use 'view assess' to see detailed results.") + + except Exception as e: + print(f"Assessment failed: {e}") + + def cmd_corr(self, args: List[str]) -> None: + """Perform correlation-based assessment.""" + if self.current_data is None: + print("No data loaded. Use 'load ' first.") + return + + if args: + try: + block_size = int(args[0]) + block_stride = int(args[1]) if len(args) > 1 else block_size + except ValueError: + print("Usage: corr [block_size] [block_stride]") + return + else: + block_size = self.settings['rndspec_options']['bss'][0] + block_stride = self.settings['rndspec_options']['bss'][1] + + print(f"Correlation assessment: block=[{block_size}, {block_stride}]") + + try: + corr_result = correlation_assessment( + time_vector=None, + signal_array=self.current_data, + block_parameters=(block_size, block_stride) + ) + + # Store result + self.data['correlation_result'] = corr_result + + # Show summary + corr_matrix = corr_result.correlation_matrix + off_diag = corr_matrix[np.triu_indices_from(corr_matrix, k=1)] + + print("Correlation assessment complete:") + print(f" Mean correlation: {off_diag.mean():.3f}") + print(f" Max correlation: {off_diag.max():.3f}") + print(f" Min correlation: {off_diag.min():.3f}") + + except Exception as e: + print(f"Correlation assessment failed: {e}") + + def cmd_rndspec(self, args: List[str]) -> None: + """Perform random projection spectral analysis.""" + if self.current_data is None: + print("No data loaded. Use 'load ' first.") + return + + # Parse parameters or use defaults + if args: + try: + reduced_dim = int(args[0]) + future = int(args[1]) if len(args) > 1 else self.settings['rndspec_options']['rfpn'][1] + past = int(args[2]) if len(args) > 2 else self.settings['rndspec_options']['rfpn'][2] + order1 = int(args[3]) if len(args) > 3 else self.settings['rndspec_options']['rfpn'][3] + order2 = int(args[4]) if len(args) > 4 else self.settings['rndspec_options']['rfpn'][4] + except ValueError: + print("Usage: rndspec [reduced_dim] [future] [past] [order1] [order2]") + return + else: + reduced_dim = self.settings['rndspec_options']['rfpn'][0] + future = self.settings['rndspec_options']['rfpn'][1] + past = self.settings['rndspec_options']['rfpn'][2] + order1 = self.settings['rndspec_options']['rfpn'][3] + order2 = self.settings['rndspec_options']['rfpn'][4] + + # Determine if SSI or AR/PCA based on parameters + if len(args) >= 3 or future != past: + # SSI mode + reduced_dimension = [reduced_dim, future, past, order1, order2] + use_cca = False + else: + # AR/PCA mode + reduced_dimension = [reduced_dim, past, order1, order2] + use_cca = False + + block_parameters = tuple(self.settings['rndspec_options']['bss']) + threshold_parameters = tuple(self.settings['rndspec_options']['thresh']) + + print(f"Random projection analysis: reduced_dim={reduced_dimension}") + print(f"Block parameters: {block_parameters}") + + try: + self.analysis_results = random_projection_spectral_analysis( + time_array=None, + signal_array=self.current_data, + block_parameters=block_parameters, + reduced_dimension=reduced_dimension, + threshold_parameters=threshold_parameters, + use_canonical_correlation_analysis=use_cca + ) + + n_blocks = len(self.analysis_results.block_results) + total_time = self.analysis_results.total_processing_time + + print(f"Analysis complete: {n_blocks} blocks processed in {total_time:.2f}s") + print("Use 'view results' to examine the results.") + + except Exception as e: + print(f"Spectral analysis failed: {e}") + + def cmd_view(self, args: List[str]) -> None: + """View analysis results.""" + if not args: + print("Usage: view ") + return + + view_type = args[0].lower() + + if view_type == 'results': + if self.analysis_results is None: + print("No spectral analysis results available.") + return + self._view_spectral_results() + + elif view_type == 'assess': + if self.assessment_results is None: + print("No assessment results available.") + return + self._view_assessment_results() + + elif view_type == 'corr': + if 'correlation_result' not in self.data: + print("No correlation results available.") + return + self._view_correlation_results() + + else: + print(f"Unknown view type: {view_type}") + + def _view_spectral_results(self) -> None: + """Display spectral analysis results summary.""" + results = self.analysis_results + print("=== Spectral Analysis Results ===") + print(f"Analysis type: {results.analysis_name}") + print(f"Total blocks: {len(results.block_results)}") + print(f"Processing time: {results.total_processing_time:.2f}s") + print(f"Block parameters: {results.block_parameters}") + print(f"Reduced dimension: {results.reduced_dimension}") + print(f"Thresholds: {results.threshold_parameters}") + + # Summary of modes found + total_modes = 0 + for block_result in results.block_results: + if hasattr(block_result.reduced_dimension_matrix, 'entries'): + total_modes += len(block_result.reduced_dimension_matrix.entries) + + print(f"Total modes identified: {total_modes}") + + def _view_assessment_results(self) -> None: + """Display assessment results summary.""" + result = self.assessment_results + print("=== AR Assessment Results ===") + print(f"Channels: {len(result.channel_names)}") + print(f"AR lag: {result.ar_lag}") + print(f"Block parameters: {result.block_parameters}") + + print("\nChannel Summary (median values):") + print("Ch# Predictability Participation RMS") + print("-" * 45) + + for i, (pred, part, rms) in enumerate(result.median_scores): + ch_name = result.channel_names[i].split('(')[0].strip() + print(f"{i+1:3d} {pred:12.3f} {part:12.3f} {rms:8.3e}") + + def _view_correlation_results(self) -> None: + """Display correlation results summary.""" + result = self.data['correlation_result'] + corr_matrix = result.correlation_matrix + + print("=== Correlation Assessment Results ===") + print(f"Channels: {len(result.channel_names)}") + + # Show correlation matrix + n_channels = len(result.channel_names) + if n_channels <= 10: + print("\nCorrelation Matrix:") + print(" ", end="") + for i in range(n_channels): + print(f"{i+1:6d}", end="") + print() + + for i in range(n_channels): + print(f"{i+1:3d}", end="") + for j in range(n_channels): + print(f"{corr_matrix[i,j]:6.2f}", end="") + print() + else: + print("(Correlation matrix too large to display)") + + def cmd_export(self, args: List[str]) -> None: + """Export analysis results.""" + if len(args) < 2: + print("Usage: export ") + print("Formats: csv, npy, npz") + return + + format_type = args[0].lower() + filename = args[1] + + try: + if format_type == 'csv': + if self.assessment_results is not None: + # Export assessment results as CSV + data = np.column_stack([ + np.arange(1, len(self.assessment_results.channel_names) + 1), + self.assessment_results.median_scores + ]) + header = "Channel,Predictability,Participation,RMS" + np.savetxt(filename, data, delimiter=',', header=header, fmt='%g') + print(f"Assessment results exported to {filename}") + else: + print("No assessment results to export") + + elif format_type in ['npy', 'npz']: + # Export all available data + export_data = {} + + if self.current_data is not None: + export_data['raw_data'] = self.current_data + + if self.assessment_results is not None: + export_data['assessment_median'] = self.assessment_results.median_scores + export_data['assessment_full'] = { + 'predictability': self.assessment_results.predictability, + 'participation': self.assessment_results.participation, + 'rms': self.assessment_results.rms_values + } + + if format_type == 'npy' and len(export_data) == 1: + np.save(filename, next(iter(export_data.values()))) + else: + np.savez(filename, **export_data) + + print(f"Data exported to {filename}") + + else: + print(f"Unsupported format: {format_type}") + + except Exception as e: + print(f"Export failed: {e}") + + def cmd_script(self, args: List[str]) -> None: + """Execute commands from script file.""" + if not args: + print("Usage: script ") + return + + script_file = args[0] + + try: + with open(script_file, 'r') as f: + commands = [line.strip() for line in f if line.strip() and not line.startswith('//')] + + self.script_commands = commands + self.script_index = 0 + self.script_mode = True + + print(f"Loaded script: {script_file} ({len(commands)} commands)") + + except Exception as e: + print(f"Failed to load script {script_file}: {e}") + + def cmd_clear(self, args: List[str]) -> None: + """Clear loaded data and results.""" + self.current_data = None + self.analysis_results = None + self.assessment_results = None + self.cluster_results = None + self.data.clear() + print("All data and results cleared.") + + +def main() -> None: + """Main entry point for eigspec CLI.""" + parser = argparse.ArgumentParser( + description="eigspec - Spectral analysis and modal identification", + formatter_class=argparse.RawDescriptionHelpFormatter + ) + + parser.add_argument( + 'script', + nargs='?', + help='Optional script file to execute' + ) + + parser.add_argument( + '--version', + action='version', + version=f'eigspec {__version__}' + ) + + parser.add_argument( + '--debug', + action='store_true', + help='Enable debug mode with detailed error messages' + ) + + args = parser.parse_args() + + if args.debug: + os.environ['EIGSPEC_DEBUG'] = '1' + + cli = EigspecCLI() + cli.run(script_file=args.script) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/tokeye/eigspec/io/__init__.py b/src/tokeye/eigspec/io/__init__.py new file mode 100644 index 0000000..66b93eb --- /dev/null +++ b/src/tokeye/eigspec/io/__init__.py @@ -0,0 +1,115 @@ +""" +Input/Output module for eigspec package. + +This module provides comprehensive data I/O functionality for scientific data formats: +- Data loading from multiple formats (HDF5, NetCDF, MATLAB, ASCII, CSV) +- Data export and saving with format validation +- Data structure definitions for time series and analysis results +- Format detection and compatibility checking +- Data validation and processing utilities + +Based on the MATLAB eigspec toolbox I/O functions: +- fetch_mpi_arrays_ptdata64.m - MPI array data loading +- bundle_shot_data.m - Shot data bundling and organization +- savebdotimage.m - B-dot image saving and export +- eigspec_mmain.m - Main I/O workflow coordination +- collect_rep_data.m - Analysis result data collection +- Various file format handlers throughout the MATLAB codebase +""" + +from .data_structures import ( + ShotData, + ProcessedData, + AnalysisResult, + TimeSeriesData, + MirnovData, + ConfigData +) + +from .loaders import ( + load_shot_data, + load_time_series, + load_mirnov_data, + load_config +) + +from .exporters import ( + save_shot_data, + save_analysis_results, + save_time_series, + export_ascii, + export_prototypes, + save_config +) + +from .formats import ( + detect_format, + supported_formats, + validate_format_support +) + +from .utils import ( + validate_data, + merge_datasets, + filter_bad_channels, + bundle_shot_data +) + +# Optional dependencies with graceful fallback +try: + import h5py + HAS_HDF5 = True +except ImportError: + HAS_HDF5 = False + +try: + import netCDF4 + HAS_NETCDF = True +except ImportError: + HAS_NETCDF = False + +try: + import scipy.io + HAS_SCIPY_IO = True +except ImportError: + HAS_SCIPY_IO = False + +__all__ = [ + # Data structures + 'ShotData', + 'ProcessedData', + 'AnalysisResult', + 'TimeSeriesData', + 'MirnovData', + 'ConfigData', + + # Loading functions + 'load_shot_data', + 'load_time_series', + 'load_mirnov_data', + 'load_config', + + # Export functions + 'save_shot_data', + 'save_analysis_results', + 'save_time_series', + 'export_ascii', + 'export_prototypes', + 'save_config', + + # Format utilities + 'detect_format', + 'supported_formats', + 'validate_format_support', + + # Data utilities + 'validate_data', + 'merge_datasets', + 'filter_bad_channels', + 'bundle_shot_data', + + # Feature flags + 'HAS_HDF5', + 'HAS_NETCDF', + 'HAS_SCIPY_IO', +] \ No newline at end of file diff --git a/src/tokeye/eigspec/io/data_structures.py b/src/tokeye/eigspec/io/data_structures.py new file mode 100644 index 0000000..69edb58 --- /dev/null +++ b/src/tokeye/eigspec/io/data_structures.py @@ -0,0 +1,298 @@ +""" +Data structures for eigspec I/O operations. + +This module defines the core data structures used for storing and managing +spectral analysis and modal identification data, including time series, +analysis results, and configuration information. +""" + +from typing import Optional, Dict, List, Any, Union, Tuple +import numpy as np +import numpy.typing as npt +from dataclasses import dataclass, field +from datetime import datetime + + +@dataclass +class TimeSeriesData: + """Time series data container. + + Attributes: + time: Time vector (seconds) + data: Data matrix (time x channels) + channels: Channel names/identifiers + sample_rate: Sampling rate (Hz) + units: Data units (e.g., 'Tesla', 'Gauss') + coordinates: Physical coordinates for each channel + metadata: Additional metadata dictionary + """ + time: npt.NDArray[np.floating] + data: npt.NDArray[np.floating] + channels: List[str] + sample_rate: float + units: str = "Tesla" + coordinates: Optional[npt.NDArray[np.floating]] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + """Validate data consistency after initialization.""" + if len(self.time) != self.data.shape[0]: + raise ValueError("Time vector length must match data rows") + if len(self.channels) != self.data.shape[1]: + raise ValueError("Number of channels must match data columns") + if self.coordinates is not None and len(self.coordinates) != len(self.channels): + raise ValueError("Coordinates length must match number of channels") + + +@dataclass +class MirnovData(TimeSeriesData): + """Mirnov probe data container (specialized time series). + + Attributes: + shot_number: Plasma shot identifier + probe_type: Type of probe ('poloidal', 'toroidal', 'mixed') + array_geometry: Geometric arrangement information + bad_channels: List of channels to exclude from analysis + calibration: Calibration factors for each channel + """ + shot_number: int = 0 + probe_type: str = "mixed" + array_geometry: Optional[Dict[str, Any]] = None + bad_channels: List[str] = field(default_factory=list) + calibration: Optional[npt.NDArray[np.floating]] = None + + +@dataclass +class ShotData: + """Plasma shot data container. + + Attributes: + shot_number: Unique shot identifier + time_range: Time range [start, end] in seconds + channels: Raw time series data organized by channel + coordinates: Physical coordinates (R, Z, phi) for each channel + channel_names: Names/identifiers for each channel + acquisition_info: Data acquisition metadata + preprocessing: Applied preprocessing steps + quality_flags: Data quality indicators + """ + shot_number: int + time_range: Tuple[float, float] + channels: Dict[str, npt.NDArray[np.floating]] + coordinates: Dict[str, Tuple[float, float, float]] + channel_names: List[str] + acquisition_info: Dict[str, Any] = field(default_factory=dict) + preprocessing: List[str] = field(default_factory=list) + quality_flags: Dict[str, bool] = field(default_factory=dict) + + def get_time_series(self, channel_subset: Optional[List[str]] = None) -> TimeSeriesData: + """Extract time series data for specified channels. + + Args: + channel_subset: List of channels to extract (None for all) + + Returns: + TimeSeriesData object with requested channels + """ + if channel_subset is None: + channel_subset = self.channel_names + + # Filter channels that exist in the data + valid_channels = [ch for ch in channel_subset if ch in self.channels] + + if not valid_channels: + raise ValueError("No valid channels found") + + # Assume all channels have the same time base (first channel) + time_vec = self.channels[valid_channels[0]][:, 0] + data_matrix = np.column_stack([self.channels[ch][:, 1] for ch in valid_channels]) + + # Extract coordinates if available + coords = None + if self.coordinates: + coords = np.array([self.coordinates.get(ch, (0, 0, 0)) for ch in valid_channels]) + + return TimeSeriesData( + time=time_vec, + data=data_matrix, + channels=valid_channels, + sample_rate=float(1.0 / np.mean(np.diff(time_vec))), + coordinates=coords, + metadata={"shot_number": self.shot_number} + ) + + +@dataclass +class ProcessedData: + """Processed/bundled data container. + + Attributes: + shot_number: Source shot identifier + time: Common time vector + data: Data matrix (time x channels) with bad channels removed + coordinates: Physical coordinates for remaining channels + channel_names: Names of remaining channels + bad_channels: List of removed channel names + processing_steps: Applied processing operations + interpolation_method: Method used for time alignment + detrend_method: Detrending method applied + filter_info: Filtering information if applied + """ + shot_number: int + time: npt.NDArray[np.floating] + data: npt.NDArray[np.floating] + coordinates: npt.NDArray[np.floating] + channel_names: List[str] + bad_channels: List[str] = field(default_factory=list) + processing_steps: List[str] = field(default_factory=list) + interpolation_method: str = "linear" + detrend_method: str = "linear" + filter_info: Optional[Dict[str, Any]] = None + + @property + def num_channels(self) -> int: + """Number of channels in processed data.""" + return len(self.channel_names) + + @property + def num_samples(self) -> int: + """Number of time samples.""" + return len(self.time) + + @property + def time_range(self) -> Tuple[float, float]: + """Time range [start, end] in seconds.""" + return (float(self.time[0]), float(self.time[-1])) + + +@dataclass +class AnalysisResult: + """Analysis results container. + + Attributes: + analysis_type: Type of analysis performed ('ssi', 'arpca', 'spectral') + shot_number: Source shot identifier + time_blocks: Time block information for analysis + frequencies: Identified frequencies (Hz) + eigenvalues: System eigenvalues (complex) + mode_shapes: Spatial mode shapes (complex) + stability: Stability indicators (damping ratios) + parameters: Analysis parameters used + quality_metrics: Quality assessment metrics + cluster_info: Clustering analysis results if applicable + timestamp: Analysis timestamp + """ + analysis_type: str + shot_number: int + time_blocks: List[Tuple[float, float]] + frequencies: npt.NDArray[np.floating] + eigenvalues: npt.NDArray[np.complexfloating] + mode_shapes: npt.NDArray[np.complexfloating] + stability: npt.NDArray[np.floating] + parameters: Dict[str, Any] + quality_metrics: Dict[str, float] = field(default_factory=dict) + cluster_info: Optional[Dict[str, Any]] = None + timestamp: datetime = field(default_factory=datetime.now) + + @property + def num_modes(self) -> int: + """Number of identified modes.""" + return len(self.frequencies) + + @property + def stable_modes(self) -> npt.NDArray[np.bool_]: + """Boolean array indicating stable modes (damping > 0).""" + return self.stability > 0 + + def get_mode_info(self, mode_idx: int) -> Dict[str, Any]: + """Get information for a specific mode. + + Args: + mode_idx: Mode index + + Returns: + Dictionary with mode information + """ + if mode_idx < 0 or mode_idx >= self.num_modes: + raise IndexError(f"Mode index {mode_idx} out of range") + + return { + "frequency": float(self.frequencies[mode_idx]), + "eigenvalue": complex(self.eigenvalues[mode_idx]), + "mode_shape": self.mode_shapes[:, mode_idx], + "stability": float(self.stability[mode_idx]), + "is_stable": bool(self.stable_modes[mode_idx]) + } + + +@dataclass +class ConfigData: + """Configuration data container. + + Attributes: + analysis_params: Analysis algorithm parameters + processing_params: Data processing parameters + plot_params: Plotting and visualization parameters + io_params: Input/output parameters + file_paths: Important file paths + user_settings: User-specific settings + version: Configuration version + """ + analysis_params: Dict[str, Any] = field(default_factory=dict) + processing_params: Dict[str, Any] = field(default_factory=dict) + plot_params: Dict[str, Any] = field(default_factory=dict) + io_params: Dict[str, Any] = field(default_factory=dict) + file_paths: Dict[str, str] = field(default_factory=dict) + user_settings: Dict[str, Any] = field(default_factory=dict) + version: str = "1.0" + + def get_param(self, category: str, param_name: str, default: Any = None) -> Any: + """Get a specific parameter value. + + Args: + category: Parameter category ('analysis', 'processing', 'plot', 'io') + param_name: Parameter name + default: Default value if parameter not found + + Returns: + Parameter value or default + """ + category_map = { + "analysis": self.analysis_params, + "processing": self.processing_params, + "plot": self.plot_params, + "io": self.io_params + } + + if category not in category_map: + raise ValueError(f"Unknown parameter category: {category}") + + return category_map[category].get(param_name, default) + + def set_param(self, category: str, param_name: str, value: Any) -> None: + """Set a specific parameter value. + + Args: + category: Parameter category + param_name: Parameter name + value: Parameter value + """ + category_map = { + "analysis": self.analysis_params, + "processing": self.processing_params, + "plot": self.plot_params, + "io": self.io_params + } + + if category not in category_map: + raise ValueError(f"Unknown parameter category: {category}") + + category_map[category][param_name] = value + + +# Type aliases for common data types +DataMatrix = npt.NDArray[np.floating] +ComplexDataMatrix = npt.NDArray[np.complexfloating] +TimeVector = npt.NDArray[np.floating] +ChannelList = List[str] +CoordinateArray = npt.NDArray[np.floating] \ No newline at end of file diff --git a/src/tokeye/eigspec/io/exporters.py b/src/tokeye/eigspec/io/exporters.py new file mode 100644 index 0000000..ddf1d3a --- /dev/null +++ b/src/tokeye/eigspec/io/exporters.py @@ -0,0 +1,694 @@ +""" +Data export functions for eigspec. + +This module provides functions to export/save data in various formats +commonly used in plasma physics and spectral analysis, including analysis +results, time series data, and configuration files. +""" + +from typing import Optional, Dict, List, Any, Union +import numpy as np +import numpy.typing as npt +from pathlib import Path +import json +import pickle +import warnings + +from .data_structures import ( + ShotData, + ProcessedData, + AnalysisResult, + TimeSeriesData, + MirnovData, + ConfigData +) + +# Optional imports with graceful fallback +try: + import h5py + HAS_HDF5 = True +except ImportError: + HAS_HDF5 = False + +try: + import netCDF4 + HAS_NETCDF = True +except ImportError: + HAS_NETCDF = False + +try: + import scipy.io as sio + HAS_SCIPY_IO = True +except ImportError: + HAS_SCIPY_IO = False + +try: + import pandas as pd + HAS_PANDAS = True +except ImportError: + HAS_PANDAS = False + + +def save_shot_data( + data: ShotData, + file_path: Union[str, Path], + file_format: str = "hdf5", + **kwargs +) -> None: + """Save shot data to file. + + Args: + data: ShotData object to save + file_path: Output file path + file_format: File format ('hdf5', 'netcdf', 'matlab', 'pickle') + **kwargs: Additional format-specific parameters + + Raises: + ValueError: If file format is not supported + ImportError: If required dependencies are missing + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if file_format == "hdf5": + _save_shot_data_hdf5(data, file_path, **kwargs) + elif file_format == "netcdf": + _save_shot_data_netcdf(data, file_path, **kwargs) + elif file_format == "matlab": + _save_shot_data_matlab(data, file_path, **kwargs) + elif file_format == "pickle": + _save_shot_data_pickle(data, file_path, **kwargs) + else: + raise ValueError(f"Unsupported file format: {file_format}") + + +def save_analysis_results( + results: AnalysisResult, + file_path: Union[str, Path], + file_format: str = "hdf5", + **kwargs +) -> None: + """Save analysis results to file. + + Args: + results: AnalysisResult object to save + file_path: Output file path + file_format: File format ('hdf5', 'matlab', 'pickle') + **kwargs: Additional format-specific parameters + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if file_format == "hdf5": + _save_results_hdf5(results, file_path, **kwargs) + elif file_format == "matlab": + _save_results_matlab(results, file_path, **kwargs) + elif file_format == "pickle": + _save_results_pickle(results, file_path, **kwargs) + else: + raise ValueError(f"Unsupported file format: {file_format}") + + +def export_ascii( + data: Union[TimeSeriesData, ProcessedData, npt.NDArray], + file_path: Union[str, Path], + header: Optional[str] = None, + delimiter: str = " ", + **kwargs +) -> None: + """Export data to ASCII format. + + Args: + data: Data to export (TimeSeriesData, ProcessedData, or numpy array) + file_path: Output file path + header: Optional header text + delimiter: Column delimiter + **kwargs: Additional parameters for numpy.savetxt + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if isinstance(data, (TimeSeriesData, ProcessedData)): + # Export time series or processed data + if hasattr(data, 'time'): + export_data = np.column_stack([data.time, data.data]) + else: + export_data = data.data + + if header is None: + if isinstance(data, TimeSeriesData) and hasattr(data, 'channels'): + header = f"Time\t{delimiter.join(data.channels)}" + elif isinstance(data, ProcessedData) and hasattr(data, 'channel_names'): + header = f"Time\t{delimiter.join(data.channel_names)}" + + elif isinstance(data, np.ndarray): + export_data = data + else: + raise ValueError("Unsupported data type for ASCII export") + + # Save using numpy.savetxt + np.savetxt( + file_path, + export_data, + delimiter=delimiter, + header=header or "", + comments='# ', + **kwargs + ) + + +def export_prototypes( + mode_shapes: npt.NDArray[np.complexfloating], + coordinates: npt.NDArray[np.floating], + file_path: Union[str, Path], + threshold: float = 0.95, + m_max: int = 8, + n_max: int = 5, + **kwargs +) -> None: + """Export mode shape prototypes in ASCII format. + + Similar to MATLAB export_prototypes_mn1_ascii function. + + Args: + mode_shapes: Complex mode shapes (channels x modes) + coordinates: Channel coordinates + file_path: Output file path + threshold: Threshold for prototype selection + m_max: Maximum poloidal mode number + n_max: Maximum toroidal mode number + **kwargs: Additional parameters + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + num_channels, num_modes = mode_shapes.shape + + # Process each mode shape + export_data = [] + + for mode_idx in range(num_modes): + mode_shape = mode_shapes[:, mode_idx] + + # Apply threshold and processing + amplitude = np.abs(mode_shape) + phase = np.angle(mode_shape) + + # Select channels above threshold + significant_channels = amplitude > threshold * np.max(amplitude) + + if np.any(significant_channels): + for ch_idx in np.where(significant_channels)[0]: + coord = coordinates[ch_idx] if coordinates is not None else [0, 0, 0] + + export_data.append([ + mode_idx + 1, # Mode number (1-indexed) + ch_idx + 1, # Channel number (1-indexed) + coord[0], # R coordinate + coord[1], # Z coordinate + coord[2], # Phi coordinate + amplitude[ch_idx], + phase[ch_idx] + ]) + + if export_data: + export_array = np.array(export_data) + header = "Mode Channel R Z Phi Amplitude Phase" + + np.savetxt( + file_path, + export_array, + header=header, + fmt=['%d', '%d', '%.6f', '%.6f', '%.6f', '%.6e', '%.6f'], + delimiter=' ', + comments='# ' + ) + else: + # Create empty file + with open(file_path, 'w') as f: + f.write("# No prototypes found above threshold\n") + + +def save_config( + config: ConfigData, + file_path: Union[str, Path], + file_format: str = "json", + **kwargs +) -> None: + """Save configuration data to file. + + Args: + config: ConfigData object to save + file_path: Output file path + file_format: File format ('json', 'matlab', 'pickle') + **kwargs: Additional format-specific parameters + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if file_format == "json": + _save_config_json(config, file_path, **kwargs) + elif file_format == "matlab": + _save_config_matlab(config, file_path, **kwargs) + elif file_format == "pickle": + _save_config_pickle(config, file_path, **kwargs) + else: + raise ValueError(f"Unsupported config format: {file_format}") + + +def save_time_series( + data: TimeSeriesData, + file_path: Union[str, Path], + file_format: str = "hdf5", + **kwargs +) -> None: + """Save time series data to file. + + Args: + data: TimeSeriesData object to save + file_path: Output file path + file_format: File format ('hdf5', 'csv', 'ascii', 'matlab') + **kwargs: Additional format-specific parameters + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if file_format == "hdf5": + _save_time_series_hdf5(data, file_path, **kwargs) + elif file_format == "csv": + _save_time_series_csv(data, file_path, **kwargs) + elif file_format == "ascii": + export_ascii(data, file_path, **kwargs) + elif file_format == "matlab": + _save_time_series_matlab(data, file_path, **kwargs) + else: + raise ValueError(f"Unsupported time series format: {file_format}") + + +def save_bdot_image( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + coordinates: Optional[npt.NDArray[np.floating]], + channel_names: List[str], + file_path: Union[str, Path], + file_format: str = "matlab" +) -> None: + """Save B-dot data in organized format. + + Similar to MATLAB savebdotimage function. + + Args: + time: Time vector + data: Data matrix (time x channels) + coordinates: Channel coordinates (channels x 3) + channel_names: Channel names + file_path: Output file path + file_format: File format ('matlab', 'hdf5') + """ + file_path = Path(file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + if coordinates is not None: + # Sort channels by array type and position + sorted_data, sorted_coords, sorted_names = _sort_array_channels( + data, coordinates, channel_names + ) + else: + sorted_data = data + sorted_coords = coordinates + sorted_names = channel_names + + if file_format == "matlab": + _save_bdot_matlab(time, sorted_data, sorted_coords, sorted_names, file_path) + elif file_format == "hdf5": + _save_bdot_hdf5(time, sorted_data, sorted_coords, sorted_names, file_path) + else: + raise ValueError(f"Unsupported format for B-dot data: {file_format}") + + +# Format-specific save functions +def _save_shot_data_hdf5(data: ShotData, file_path: Path, **kwargs) -> None: + """Save shot data to HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + import h5py + + with h5py.File(file_path, 'w') as f: + # Save metadata as attributes + f.attrs['shot_number'] = data.shot_number + f.attrs['time_range'] = data.time_range + + # Save channel data + if data.channels: + channels_group = f.create_group('channels') + for ch_name, ch_data in data.channels.items(): + ch_dataset = channels_group.create_dataset(ch_name, data=ch_data) + + # Save coordinates as attributes + if ch_name in data.coordinates: + ch_dataset.attrs['coordinates'] = data.coordinates[ch_name] + + # Save channel names + if data.channel_names: + f.create_dataset('channel_names', data=[s.encode() for s in data.channel_names]) + + # Save acquisition info + for key, value in data.acquisition_info.items(): + if isinstance(value, (int, float, str)): + f.attrs[f'acq_{key}'] = value + + +def _save_shot_data_matlab(data: ShotData, file_path: Path, **kwargs) -> None: + """Save shot data to MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + import scipy.io as sio + + # Prepare data for MATLAB format + save_dict = { + 'shot_number': data.shot_number, + 'time_range': data.time_range, + 'channel_names': data.channel_names + } + + # Concatenate channel data if available + if data.channels: + channel_data = [] + for ch_name in data.channel_names: + if ch_name in data.channels: + channel_data.append(data.channels[ch_name]) + + if channel_data: + save_dict['channels'] = np.array(channel_data) + + # Save coordinates + if data.coordinates: + coords_array = np.array([data.coordinates.get(ch, [0, 0, 0]) for ch in data.channel_names]) + save_dict['coordinates'] = coords_array + + sio.savemat(file_path, save_dict) + + +def _save_shot_data_netcdf(data: ShotData, file_path: Path, **kwargs) -> None: + """Save shot data to NetCDF file.""" + if not HAS_NETCDF: + raise ImportError("netCDF4 is required for NetCDF support") + + import netCDF4 + + with netCDF4.Dataset(file_path, 'w') as nc: + # Global attributes + nc.setncattr('shot_number', data.shot_number) + nc.setncattr('time_range_start', data.time_range[0]) + nc.setncattr('time_range_end', data.time_range[1]) + + # Create dimensions + if data.channels: + # Determine dimensions from first channel + first_channel = next(iter(data.channels.values())) + time_dim = nc.createDimension('time', first_channel.shape[0]) + channel_dim = nc.createDimension('channels', len(data.channel_names)) + + # Create time variable + time_var = nc.createVariable('time', 'f8', ('time',)) + time_var[:] = first_channel[:, 0] + time_var.units = 'seconds' + + # Create channel variables + for i, ch_name in enumerate(data.channel_names): + if ch_name in data.channels: + ch_var = nc.createVariable(f'channel_{i:03d}', 'f8', ('time',)) + ch_var[:] = data.channels[ch_name][:, 1] + ch_var.channel_name = ch_name + + +def _save_shot_data_pickle(data: ShotData, file_path: Path, **kwargs) -> None: + """Save shot data to pickle file.""" + with open(file_path, 'wb') as f: + pickle.dump(data, f) + + +def _save_results_hdf5(results: AnalysisResult, file_path: Path, **kwargs) -> None: + """Save analysis results to HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + import h5py + + with h5py.File(file_path, 'w') as f: + # Save basic attributes + f.attrs['analysis_type'] = results.analysis_type.encode() + f.attrs['shot_number'] = results.shot_number + f.attrs['timestamp'] = results.timestamp.isoformat().encode() + + # Save datasets + f.create_dataset('time_blocks', data=results.time_blocks) + f.create_dataset('frequencies', data=results.frequencies) + f.create_dataset('eigenvalues', data=results.eigenvalues) + f.create_dataset('mode_shapes', data=results.mode_shapes) + f.create_dataset('stability', data=results.stability) + + # Save parameters as attributes + for key, value in results.parameters.items(): + if isinstance(value, (int, float, str)): + f.attrs[f'param_{key}'] = value + + +def _save_results_matlab(results: AnalysisResult, file_path: Path, **kwargs) -> None: + """Save analysis results to MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + import scipy.io as sio + + save_dict = { + 'analysis_type': results.analysis_type, + 'shot_number': results.shot_number, + 'time_blocks': results.time_blocks, + 'frequencies': results.frequencies, + 'eigenvalues': results.eigenvalues, + 'mode_shapes': results.mode_shapes, + 'stability': results.stability, + 'parameters': results.parameters, + 'timestamp': results.timestamp.isoformat() + } + + sio.savemat(file_path, save_dict) + + +def _save_results_pickle(results: AnalysisResult, file_path: Path, **kwargs) -> None: + """Save analysis results to pickle file.""" + with open(file_path, 'wb') as f: + pickle.dump(results, f) + + +def _save_config_json(config: ConfigData, file_path: Path, **kwargs) -> None: + """Save configuration to JSON file.""" + config_dict = { + 'analysis_params': config.analysis_params, + 'processing_params': config.processing_params, + 'plot_params': config.plot_params, + 'io_params': config.io_params, + 'file_paths': config.file_paths, + 'user_settings': config.user_settings, + 'version': config.version + } + + with open(file_path, 'w') as f: + json.dump(config_dict, f, indent=2, default=_json_serializer) + + +def _save_config_matlab(config: ConfigData, file_path: Path, **kwargs) -> None: + """Save configuration to MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + import scipy.io as sio + + save_dict = { + 'analysis_params': config.analysis_params, + 'processing_params': config.processing_params, + 'plot_params': config.plot_params, + 'io_params': config.io_params, + 'file_paths': config.file_paths, + 'user_settings': config.user_settings, + 'version': config.version + } + + sio.savemat(file_path, save_dict) + + +def _save_config_pickle(config: ConfigData, file_path: Path, **kwargs) -> None: + """Save configuration to pickle file.""" + with open(file_path, 'wb') as f: + pickle.dump(config, f) + + +def _save_time_series_hdf5(data: TimeSeriesData, file_path: Path, **kwargs) -> None: + """Save time series to HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + import h5py + + with h5py.File(file_path, 'w') as f: + f.create_dataset('time', data=data.time) + f.create_dataset('data', data=data.data) + f.create_dataset('channel_names', data=[s.encode() for s in data.channels]) + + f.attrs['sample_rate'] = data.sample_rate + f.attrs['units'] = data.units.encode() + + if data.coordinates is not None: + f.create_dataset('coordinates', data=data.coordinates) + + # Save metadata + for key, value in data.metadata.items(): + if isinstance(value, (int, float, str)): + f.attrs[f'meta_{key}'] = value + + +def _save_time_series_csv(data: TimeSeriesData, file_path: Path, **kwargs) -> None: + """Save time series to CSV file.""" + if HAS_PANDAS: + import pandas as pd + + # Create DataFrame + df_data = {'time': data.time} + for i, ch_name in enumerate(data.channels): + df_data[ch_name] = data.data[:, i] + + df = pd.DataFrame(df_data) + df.to_csv(file_path, index=False, **kwargs) + else: + # Fallback to ASCII export + export_ascii(data, file_path, delimiter=',', **kwargs) + + +def _save_time_series_matlab(data: TimeSeriesData, file_path: Path, **kwargs) -> None: + """Save time series to MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + import scipy.io as sio + + save_dict = { + 'T': data.time, + 'Y': data.data, + 'ynames': data.channels, + 'sample_rate': data.sample_rate, + 'units': data.units + } + + if data.coordinates is not None: + save_dict['Yxy'] = data.coordinates + + sio.savemat(file_path, save_dict) + + +def _save_bdot_matlab( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + coordinates: Optional[npt.NDArray[np.floating]], + channel_names: List[str], + file_path: Path +) -> None: + """Save B-dot data to MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + import scipy.io as sio + + save_dict = { + 'T': time, + 'Y': data, + 'ynames': channel_names + } + + if coordinates is not None: + save_dict['Yxy'] = coordinates + + sio.savemat(file_path, save_dict) + + +def _save_bdot_hdf5( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + coordinates: Optional[npt.NDArray[np.floating]], + channel_names: List[str], + file_path: Path +) -> None: + """Save B-dot data to HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + import h5py + + with h5py.File(file_path, 'w') as f: + f.create_dataset('time', data=time) + f.create_dataset('data', data=data) + f.create_dataset('channel_names', data=[s.encode() for s in channel_names]) + + if coordinates is not None: + f.create_dataset('coordinates', data=coordinates) + + +def _sort_array_channels( + data: npt.NDArray[np.floating], + coordinates: npt.NDArray[np.floating], + channel_names: List[str] +) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], List[str]]: + """Sort channels by array type (toroidal/poloidal) and position.""" + # Identify toroidal and poloidal arrays based on coordinates + # This is a simplified version of the MATLAB logic + eps_x = 2e-2 + eps_y = 1e-1 + pol_array_y = 5.6287 + tor_array_x = 0.0 + + # Find indices for different array types + tor_indices = np.where(np.abs(coordinates[:, 0] - tor_array_x) <= eps_x)[0] + pol_indices = np.where(np.abs(coordinates[:, 1] - pol_array_y) <= eps_y)[0] + + # Sort toroidal array by y-coordinate + if len(tor_indices) > 0: + tor_sort_idx = np.argsort(coordinates[tor_indices, 1]) + tor_indices = tor_indices[tor_sort_idx] + + # Sort poloidal array by x-coordinate + if len(pol_indices) > 0: + pol_sort_idx = np.argsort(coordinates[pol_indices, 0]) + pol_indices = pol_indices[pol_sort_idx] + + # Remove overlap between arrays + pol_indices = np.setdiff1d(pol_indices, tor_indices) + + # Get remaining indices + all_indices = np.arange(len(channel_names)) + other_indices = np.setdiff1d(all_indices, np.concatenate([tor_indices, pol_indices])) + + # Combine in order: toroidal, poloidal, others + new_order = np.concatenate([tor_indices, pol_indices, other_indices]) + + # Reorder data, coordinates, and names + sorted_data = data[:, new_order] + sorted_coords = coordinates[new_order] + sorted_names = [channel_names[i] for i in new_order] + + return sorted_data, sorted_coords, sorted_names + + +def _json_serializer(obj: Any) -> Any: + """JSON serializer for numpy arrays and other objects.""" + if isinstance(obj, np.ndarray): + return obj.tolist() + elif isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.complexfloating): + return {'real': float(obj.real), 'imag': float(obj.imag)} + else: + raise TypeError(f"Object of type {type(obj)} is not JSON serializable") \ No newline at end of file diff --git a/src/tokeye/eigspec/io/formats.py b/src/tokeye/eigspec/io/formats.py new file mode 100644 index 0000000..ebebbed --- /dev/null +++ b/src/tokeye/eigspec/io/formats.py @@ -0,0 +1,458 @@ +""" +File format detection and validation for eigspec I/O. + +This module provides utilities for detecting file formats, validating +file contents, and managing supported format information. +""" + +from typing import List, Dict, Optional, Set, Any +from pathlib import Path +import numpy as np +import numpy.typing as npt + + +# Supported file format definitions +SUPPORTED_FORMATS = { + 'hdf5': { + 'extensions': ['.h5', '.hdf5'], + 'description': 'Hierarchical Data Format 5', + 'features': ['compression', 'metadata', 'complex_data', 'large_files'], + 'dependencies': ['h5py'] + }, + 'netcdf': { + 'extensions': ['.nc', '.netcdf'], + 'description': 'Network Common Data Form', + 'features': ['metadata', 'self_describing', 'portable'], + 'dependencies': ['netCDF4'] + }, + 'matlab': { + 'extensions': ['.mat'], + 'description': 'MATLAB MAT-file', + 'features': ['matlab_compatible', 'structured_data'], + 'dependencies': ['scipy'] + }, + 'ascii': { + 'extensions': ['.txt', '.dat', '.asc'], + 'description': 'ASCII text format', + 'features': ['human_readable', 'portable', 'simple'], + 'dependencies': [] + }, + 'csv': { + 'extensions': ['.csv'], + 'description': 'Comma Separated Values', + 'features': ['human_readable', 'tabular', 'excel_compatible'], + 'dependencies': [] + }, + 'binary': { + 'extensions': ['.bin', '.raw'], + 'description': 'Binary data format', + 'features': ['compact', 'fast_io'], + 'dependencies': [] + }, + 'pickle': { + 'extensions': ['.pkl', '.pickle'], + 'description': 'Python pickle format', + 'features': ['python_objects', 'complete_serialization'], + 'dependencies': [] + }, + 'json': { + 'extensions': ['.json'], + 'description': 'JavaScript Object Notation', + 'features': ['human_readable', 'web_compatible', 'portable'], + 'dependencies': [] + } +} + +# File type signatures for format detection +FILE_SIGNATURES = { + 'hdf5': [ + b'\x89HDF\r\n\x1a\n', # HDF5 signature + ], + 'netcdf': [ + b'CDF\x01', # NetCDF classic + b'CDF\x02', # NetCDF 64-bit offset + b'\x89HDF\r\n\x1a\n\x00\x00\x00\x08\x00\x08\x00\x00', # NetCDF-4 + ], + 'matlab': [ + b'MATLAB', # MATLAB v4 format + b'\x00\x01IM', # Some MATLAB formats + ], + 'gzip': [ + b'\x1f\x8b', # Gzip compressed + ] +} + + +def detect_format(file_path: Path) -> str: + """Detect file format from path and content. + + Args: + file_path: Path to file + + Returns: + Detected format string + + Raises: + FileNotFoundError: If file does not exist + """ + if not file_path.exists(): + raise FileNotFoundError(f"File not found: {file_path}") + + # First try extension-based detection + format_from_ext = _detect_from_extension(file_path) + if format_from_ext: + return format_from_ext + + # Then try content-based detection + format_from_content = _detect_from_content(file_path) + if format_from_content: + return format_from_content + + # Default fallback + return 'ascii' + + +def supported_formats() -> Dict[str, Dict[str, Any]]: + """Get information about supported file formats. + + Returns: + Dictionary with format information + """ + return SUPPORTED_FORMATS.copy() + + +def get_format_info(format_name: str) -> Optional[Dict[str, Any]]: + """Get information about a specific format. + + Args: + format_name: Name of the format + + Returns: + Format information dictionary or None if not found + """ + return SUPPORTED_FORMATS.get(format_name) + + +def validate_format_support(format_name: str) -> tuple[bool, List[str]]: + """Check if a format is supported and what dependencies are missing. + + Args: + format_name: Name of the format to check + + Returns: + Tuple of (is_supported, missing_dependencies) + """ + if format_name not in SUPPORTED_FORMATS: + return False, [f"Unknown format: {format_name}"] + + format_info = SUPPORTED_FORMATS[format_name] + missing_deps = [] + + for dep in format_info.get('dependencies', []): + try: + if dep == 'h5py': + import h5py + elif dep == 'netCDF4': + import netCDF4 + elif dep == 'scipy': + import scipy.io + elif dep == 'pandas': + import pandas + except ImportError: + missing_deps.append(dep) + + return len(missing_deps) == 0, missing_deps + + +def get_recommended_format( + data_type: str, + features_needed: Optional[List[str]] = None +) -> str: + """Get recommended format for a given data type and features. + + Args: + data_type: Type of data ('time_series', 'analysis_results', 'config') + features_needed: List of required features + + Returns: + Recommended format name + """ + if features_needed is None: + features_needed = [] + + # Default recommendations by data type + defaults = { + 'time_series': 'hdf5', + 'analysis_results': 'hdf5', + 'config': 'json', + 'shot_data': 'hdf5', + 'processed_data': 'hdf5' + } + + if data_type in defaults: + recommended = defaults[data_type] + + # Check if recommended format supports needed features + format_info = SUPPORTED_FORMATS.get(recommended, {}) + format_features = format_info.get('features', []) + + if all(feature in format_features for feature in features_needed): + is_supported, _ = validate_format_support(recommended) + if is_supported: + return recommended + + # Fallback selection based on features + for format_name, format_info in SUPPORTED_FORMATS.items(): + format_features = format_info.get('features', []) + if all(feature in format_features for feature in features_needed): + is_supported, _ = validate_format_support(format_name) + if is_supported: + return format_name + + # Final fallback + return 'ascii' + + +def validate_file_structure( + file_path: Path, + expected_format: str, + required_fields: Optional[List[str]] = None +) -> tuple[bool, List[str]]: + """Validate file structure and content. + + Args: + file_path: Path to file to validate + expected_format: Expected file format + required_fields: List of required data fields + + Returns: + Tuple of (is_valid, error_messages) + """ + if not file_path.exists(): + return False, [f"File not found: {file_path}"] + + errors = [] + + # Check format detection matches expectation + detected_format = detect_format(file_path) + if detected_format != expected_format: + errors.append(f"Format mismatch: expected {expected_format}, detected {detected_format}") + + # Format-specific validation + if expected_format == 'hdf5': + errors.extend(_validate_hdf5_structure(file_path, required_fields)) + elif expected_format == 'netcdf': + errors.extend(_validate_netcdf_structure(file_path, required_fields)) + elif expected_format == 'matlab': + errors.extend(_validate_matlab_structure(file_path, required_fields)) + elif expected_format in ['ascii', 'csv']: + errors.extend(_validate_text_structure(file_path, required_fields)) + + return len(errors) == 0, errors + + +def get_format_extensions() -> Dict[str, List[str]]: + """Get file extensions for each supported format. + + Returns: + Dictionary mapping format names to extension lists + """ + return {name: info['extensions'] for name, info in SUPPORTED_FORMATS.items()} + + +def extension_to_format(extension: str) -> Optional[str]: + """Map file extension to format name. + + Args: + extension: File extension (with or without leading dot) + + Returns: + Format name or None if not found + """ + if not extension.startswith('.'): + extension = '.' + extension + + extension = extension.lower() + + for format_name, format_info in SUPPORTED_FORMATS.items(): + if extension in format_info['extensions']: + return format_name + + return None + + +def _detect_from_extension(file_path: Path) -> Optional[str]: + """Detect format from file extension.""" + suffix = file_path.suffix.lower() + return extension_to_format(suffix) + + +def _detect_from_content(file_path: Path) -> Optional[str]: + """Detect format from file content.""" + try: + with open(file_path, 'rb') as f: + header = f.read(32) # Read more bytes for better detection + + for format_name, signatures in FILE_SIGNATURES.items(): + for signature in signatures: + if header.startswith(signature): + return format_name + + # Try to detect text vs binary + try: + header.decode('utf-8') + return 'ascii' # Likely text file + except UnicodeDecodeError: + return 'binary' # Binary file + + except Exception: + return None + + +def _validate_hdf5_structure( + file_path: Path, + required_fields: Optional[List[str]] = None +) -> List[str]: + """Validate HDF5 file structure.""" + errors = [] + + try: + import h5py + + with h5py.File(file_path, 'r') as f: + if required_fields: + for field in required_fields: + if field not in f and field not in f.attrs: + errors.append(f"Required field '{field}' not found in HDF5 file") + + except ImportError: + errors.append("h5py not available for HDF5 validation") + except Exception as e: + errors.append(f"Error reading HDF5 file: {e}") + + return errors + + +def _validate_netcdf_structure( + file_path: Path, + required_fields: Optional[List[str]] = None +) -> List[str]: + """Validate NetCDF file structure.""" + errors = [] + + try: + import netCDF4 + + with netCDF4.Dataset(file_path, 'r') as nc: + if required_fields: + available_vars = set(nc.variables.keys()) + available_attrs = set(nc.ncattrs()) + + for field in required_fields: + if field not in available_vars and field not in available_attrs: + errors.append(f"Required field '{field}' not found in NetCDF file") + + except ImportError: + errors.append("netCDF4 not available for NetCDF validation") + except Exception as e: + errors.append(f"Error reading NetCDF file: {e}") + + return errors + + +def _validate_matlab_structure( + file_path: Path, + required_fields: Optional[List[str]] = None +) -> List[str]: + """Validate MATLAB file structure.""" + errors = [] + + try: + import scipy.io as sio + + data = sio.loadmat(str(file_path)) + + if required_fields: + available_fields = set(data.keys()) + + for field in required_fields: + if field not in available_fields: + errors.append(f"Required field '{field}' not found in MATLAB file") + + except ImportError: + errors.append("scipy not available for MATLAB validation") + except Exception as e: + errors.append(f"Error reading MATLAB file: {e}") + + return errors + + +def _validate_text_structure( + file_path: Path, + required_fields: Optional[List[str]] = None +) -> List[str]: + """Validate text file structure.""" + errors = [] + + try: + with open(file_path, 'r') as f: + lines = f.readlines() + + if not lines: + errors.append("Text file is empty") + return errors + + # Check if first line looks like a header + first_line = lines[0].strip() + if first_line.startswith('#'): + # Header line, extract field names + header_fields = first_line[1:].strip().split() + + if required_fields: + for field in required_fields: + if field not in header_fields: + errors.append(f"Required field '{field}' not found in text file header") + + # Check data consistency + if len(lines) > 1: + try: + # Try to parse a data line + data_line = lines[1 if first_line.startswith('#') else 0] + data_cols = data_line.strip().split() + + # Check if all lines have consistent number of columns + for i, line in enumerate(lines[1:], start=2): + if not line.strip() or line.strip().startswith('#'): + continue + cols = line.strip().split() + if len(cols) != len(data_cols): + errors.append(f"Inconsistent number of columns at line {i}") + break + + except Exception as e: + errors.append(f"Error parsing text file data: {e}") + + except Exception as e: + errors.append(f"Error reading text file: {e}") + + return errors + + +def format_compatibility_matrix() -> Dict[str, Dict[str, bool]]: + """Get compatibility matrix between formats. + + Returns: + Dictionary showing which formats can be converted to which others + """ + # This is a simplified compatibility matrix + # In practice, compatibility depends on the specific data structure + return { + 'hdf5': {'netcdf': True, 'matlab': True, 'ascii': True, 'csv': True, 'pickle': True}, + 'netcdf': {'hdf5': True, 'matlab': True, 'ascii': True, 'csv': True}, + 'matlab': {'hdf5': True, 'ascii': True, 'csv': True, 'pickle': True}, + 'ascii': {'hdf5': True, 'netcdf': True, 'matlab': True, 'csv': True}, + 'csv': {'hdf5': True, 'netcdf': True, 'matlab': True, 'ascii': True}, + 'pickle': {'hdf5': True, 'matlab': True, 'ascii': True}, + 'binary': {'ascii': False, 'csv': False}, # Limited conversion options + 'json': {'hdf5': True, 'matlab': True, 'ascii': True, 'pickle': True} + } \ No newline at end of file diff --git a/src/tokeye/eigspec/io/loaders.py b/src/tokeye/eigspec/io/loaders.py new file mode 100644 index 0000000..c408a5b --- /dev/null +++ b/src/tokeye/eigspec/io/loaders.py @@ -0,0 +1,702 @@ +""" +Data loading functions for eigspec. + +This module provides functions to load various data formats commonly used +in plasma physics and spectral analysis, including time series data, +shot data, configuration files, and analysis results. +""" + +from typing import Optional, Dict, List, Any, Union, Tuple +import numpy as np +import numpy.typing as npt +from pathlib import Path +import warnings +import json +import pickle + +from .data_structures import ( + ShotData, + ProcessedData, + AnalysisResult, + TimeSeriesData, + MirnovData, + ConfigData +) + +# Optional imports with graceful fallback +try: + import h5py + HAS_HDF5 = True +except ImportError: + HAS_HDF5 = False + +try: + import netCDF4 + HAS_NETCDF = True +except ImportError: + HAS_NETCDF = False + +try: + import scipy.io as sio + HAS_SCIPY_IO = True +except ImportError: + HAS_SCIPY_IO = False + +try: + import pandas as pd + HAS_PANDAS = True +except ImportError: + HAS_PANDAS = False + + +def load_shot_data( + file_path: Union[str, Path], + file_format: Optional[str] = None, + **kwargs +) -> ShotData: + """Load plasma shot data from file. + + Args: + file_path: Path to data file + file_format: File format ('hdf5', 'netcdf', 'matlab', 'ascii', None for auto-detect) + **kwargs: Additional format-specific parameters + + Returns: + ShotData object with loaded data + + Raises: + ValueError: If file format is not supported or file cannot be read + FileNotFoundError: If file does not exist + """ + file_path = Path(file_path) + if not file_path.exists(): + raise FileNotFoundError(f"File not found: {file_path}") + + if file_format is None: + file_format = _detect_file_format(file_path) + + if file_format == "hdf5": + return _load_shot_data_hdf5(file_path, **kwargs) + elif file_format == "netcdf": + return _load_shot_data_netcdf(file_path, **kwargs) + elif file_format == "matlab": + return _load_shot_data_matlab(file_path, **kwargs) + elif file_format == "ascii": + return _load_shot_data_ascii(file_path, **kwargs) + elif file_format == "binary": + return _load_shot_data_binary(file_path, **kwargs) + else: + raise ValueError(f"Unsupported file format: {file_format}") + + +def load_time_series( + file_path: Union[str, Path], + file_format: Optional[str] = None, + **kwargs +) -> TimeSeriesData: + """Load time series data from file. + + Args: + file_path: Path to data file + file_format: File format (None for auto-detect) + **kwargs: Additional format-specific parameters + + Returns: + TimeSeriesData object with loaded data + """ + file_path = Path(file_path) + if not file_path.exists(): + raise FileNotFoundError(f"File not found: {file_path}") + + if file_format is None: + file_format = _detect_file_format(file_path) + + if file_format == "ascii": + return _load_time_series_ascii(file_path, **kwargs) + elif file_format == "csv": + return _load_time_series_csv(file_path, **kwargs) + elif file_format == "hdf5": + return _load_time_series_hdf5(file_path, **kwargs) + elif file_format == "matlab": + return _load_time_series_matlab(file_path, **kwargs) + else: + raise ValueError(f"Unsupported file format: {file_format}") + + +def load_mirnov_data( + file_path: Union[str, Path], + shot_number: int, + probe_type: str = "mixed", + **kwargs +) -> MirnovData: + """Load Mirnov probe data from file. + + Args: + file_path: Path to data file + shot_number: Plasma shot number + probe_type: Type of probe array ('poloidal', 'toroidal', 'mixed') + **kwargs: Additional parameters + + Returns: + MirnovData object with loaded probe data + """ + # Load as time series first + ts_data = load_time_series(file_path, **kwargs) + + # Convert to MirnovData + return MirnovData( + time=ts_data.time, + data=ts_data.data, + channels=ts_data.channels, + sample_rate=ts_data.sample_rate, + units=ts_data.units, + coordinates=ts_data.coordinates, + metadata=ts_data.metadata, + shot_number=shot_number, + probe_type=probe_type + ) + + +def load_config( + file_path: Union[str, Path], + file_format: Optional[str] = None +) -> ConfigData: + """Load configuration data from file. + + Args: + file_path: Path to configuration file + file_format: File format ('json', 'yaml', 'matlab', None for auto-detect) + + Returns: + ConfigData object with configuration + """ + file_path = Path(file_path) + if not file_path.exists(): + raise FileNotFoundError(f"Configuration file not found: {file_path}") + + if file_format is None: + file_format = _detect_config_format(file_path) + + if file_format == "json": + return _load_config_json(file_path) + elif file_format == "matlab": + return _load_config_matlab(file_path) + elif file_format == "pickle": + return _load_config_pickle(file_path) + else: + raise ValueError(f"Unsupported config format: {file_format}") + + +def load_analysis_results( + file_path: Union[str, Path], + file_format: Optional[str] = None +) -> AnalysisResult: + """Load analysis results from file. + + Args: + file_path: Path to results file + file_format: File format (None for auto-detect) + + Returns: + AnalysisResult object with loaded results + """ + file_path = Path(file_path) + if not file_path.exists(): + raise FileNotFoundError(f"Results file not found: {file_path}") + + if file_format is None: + file_format = _detect_file_format(file_path) + + if file_format == "hdf5": + return _load_results_hdf5(file_path) + elif file_format == "matlab": + return _load_results_matlab(file_path) + elif file_format == "pickle": + return _load_results_pickle(file_path) + else: + raise ValueError(f"Unsupported results format: {file_format}") + + +# Format detection utilities +def _detect_file_format(file_path: Path) -> str: + """Detect file format from extension and content.""" + suffix = file_path.suffix.lower() + + if suffix in [".h5", ".hdf5"]: + return "hdf5" + elif suffix in [".nc", ".netcdf"]: + return "netcdf" + elif suffix in [".mat"]: + return "matlab" + elif suffix in [".txt", ".dat", ".asc"]: + return "ascii" + elif suffix in [".csv"]: + return "csv" + elif suffix in [".bin", ".raw"]: + return "binary" + elif suffix in [".pkl", ".pickle"]: + return "pickle" + else: + # Try to detect from content + return _detect_from_content(file_path) + + +def _detect_config_format(file_path: Path) -> str: + """Detect configuration file format.""" + suffix = file_path.suffix.lower() + + if suffix in [".json"]: + return "json" + elif suffix in [".mat"]: + return "matlab" + elif suffix in [".pkl", ".pickle"]: + return "pickle" + else: + return "json" # Default + + +def _detect_from_content(file_path: Path) -> str: + """Detect format from file content.""" + try: + with open(file_path, 'rb') as f: + header = f.read(16) + + # Check for HDF5 signature + if header.startswith(b'\x89HDF\r\n\x1a\n'): + return "hdf5" + + # Check for NetCDF signature + if header.startswith(b'CDF\x01') or header.startswith(b'CDF\x02'): + return "netcdf" + + # Default to ASCII + return "ascii" + except: + return "ascii" + + +# Format-specific loaders +def _load_shot_data_hdf5(file_path: Path, **kwargs) -> ShotData: + """Load shot data from HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + with h5py.File(file_path, 'r') as f: + shot_number = int(f.attrs.get('shot_number', 0)) + time_range = tuple(f.attrs.get('time_range', (0.0, 1.0))) + + # Load channel data + channels = {} + coordinates = {} + channel_names = [] + + if 'channels' in f: + for ch_name in f['channels'].keys(): + ch_data = f['channels'][ch_name][...] + channels[ch_name] = ch_data + channel_names.append(ch_name) + + # Load coordinates if available + if 'coordinates' in f['channels'][ch_name].attrs: + coords = f['channels'][ch_name].attrs['coordinates'] + coordinates[ch_name] = tuple(coords) + + # Load metadata + acquisition_info = dict(f.attrs) if f.attrs else {} + + return ShotData( + shot_number=shot_number, + time_range=time_range, + channels=channels, + coordinates=coordinates, + channel_names=channel_names, + acquisition_info=acquisition_info + ) + + +def _load_shot_data_matlab(file_path: Path, **kwargs) -> ShotData: + """Load shot data from MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + data = sio.loadmat(str(file_path)) + + # Extract shot information + shot_number = int(data.get('shot_number', [0])[0]) + time_range = tuple(data.get('time_range', [0.0, 1.0])) + + # Load channel data + channels = {} + coordinates = {} + channel_names = [] + + if 'channels' in data: + ch_data = data['channels'] + for i, ch_name in enumerate(data.get('channel_names', [])): + if isinstance(ch_name, np.ndarray): + ch_name = str(ch_name[0]) + channels[ch_name] = ch_data[:, i:i+2] # time and data + channel_names.append(ch_name) + + return ShotData( + shot_number=shot_number, + time_range=time_range, + channels=channels, + coordinates=coordinates, + channel_names=channel_names + ) + + +def _load_shot_data_ascii(file_path: Path, **kwargs) -> ShotData: + """Load shot data from ASCII file.""" + delimiter = kwargs.get('delimiter', None) + skip_header = kwargs.get('skip_header', 0) + + data = np.loadtxt(file_path, delimiter=delimiter, skiprows=skip_header) + + # Assume first column is time, rest are channels + time_vec = data[:, 0] + channels = {} + channel_names = [] + + for i in range(1, data.shape[1]): + ch_name = f"channel_{i}" + channels[ch_name] = np.column_stack([time_vec, data[:, i]]) + channel_names.append(ch_name) + + return ShotData( + shot_number=kwargs.get('shot_number', 0), + time_range=(float(time_vec[0]), float(time_vec[-1])), + channels=channels, + coordinates={}, + channel_names=channel_names + ) + + +def _load_shot_data_binary(file_path: Path, **kwargs) -> ShotData: + """Load shot data from binary file.""" + dtype = kwargs.get('dtype', np.float64) + shape = kwargs.get('shape', None) + + data = np.fromfile(file_path, dtype=dtype) + + if shape is not None: + data = data.reshape(shape) + + # Assume first column is time + time_vec = data[:, 0] + channels = {} + channel_names = [] + + for i in range(1, data.shape[1]): + ch_name = f"channel_{i}" + channels[ch_name] = np.column_stack([time_vec, data[:, i]]) + channel_names.append(ch_name) + + return ShotData( + shot_number=kwargs.get('shot_number', 0), + time_range=(float(time_vec[0]), float(time_vec[-1])), + channels=channels, + coordinates={}, + channel_names=channel_names + ) + + +def _load_shot_data_netcdf(file_path: Path, **kwargs) -> ShotData: + """Load shot data from NetCDF file.""" + if not HAS_NETCDF: + raise ImportError("netCDF4 is required for NetCDF support") + + with netCDF4.Dataset(file_path, 'r') as nc: + shot_number = int(nc.getncattr('shot_number') if 'shot_number' in nc.ncattrs() else 0) + + # Load variables + channels = {} + channel_names = [] + coordinates = {} + + for var_name in nc.variables: + if var_name != 'time': + var_data = nc.variables[var_name][...] + time_data = nc.variables['time'][...] + channels[var_name] = np.column_stack([time_data, var_data]) + channel_names.append(var_name) + + time_range = (float(time_data[0]), float(time_data[-1])) if 'time_data' in locals() else (0.0, 1.0) + + return ShotData( + shot_number=shot_number, + time_range=time_range, + channels=channels, + coordinates=coordinates, + channel_names=channel_names + ) + + +def _load_time_series_ascii(file_path: Path, **kwargs) -> TimeSeriesData: + """Load time series from ASCII file.""" + delimiter = kwargs.get('delimiter', None) + skip_header = kwargs.get('skip_header', 0) + channel_names = kwargs.get('channel_names', None) + + data = np.loadtxt(file_path, delimiter=delimiter, skiprows=skip_header) + + # Assume first column is time + time_vec = data[:, 0] + data_matrix = data[:, 1:] + + if channel_names is None: + channel_names = [f"channel_{i+1}" for i in range(data_matrix.shape[1])] + + sample_rate = float(1.0 / np.mean(np.diff(time_vec))) if len(time_vec) > 1 else 1.0 + + return TimeSeriesData( + time=time_vec, + data=data_matrix, + channels=channel_names, + sample_rate=sample_rate, + units=kwargs.get('units', 'unknown') + ) + + +def _load_time_series_csv(file_path: Path, **kwargs) -> TimeSeriesData: + """Load time series from CSV file.""" + if not HAS_PANDAS: + # Fallback to numpy + return _load_time_series_ascii(file_path, delimiter=',', **kwargs) + + df = pd.read_csv(file_path, **kwargs) + + # Assume first column is time + time_col = df.columns[0] + time_vec = df[time_col].values + + channel_names = list(df.columns[1:]) + data_matrix = df[channel_names].values + + sample_rate = float(1.0 / np.mean(np.diff(time_vec))) if len(time_vec) > 1 else 1.0 + + return TimeSeriesData( + time=time_vec, + data=data_matrix, + channels=channel_names, + sample_rate=sample_rate + ) + + +def _load_time_series_hdf5(file_path: Path, **kwargs) -> TimeSeriesData: + """Load time series from HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + with h5py.File(file_path, 'r') as f: + time_vec = f['time'][...] + data_matrix = f['data'][...] + + channel_names = [] + if 'channel_names' in f: + channel_names = [s.decode() if isinstance(s, bytes) else str(s) for s in f['channel_names'][...]] + else: + channel_names = [f"channel_{i+1}" for i in range(data_matrix.shape[1])] + + sample_rate = float(f.attrs.get('sample_rate', 1.0 / np.mean(np.diff(time_vec)))) + units = f.attrs.get('units', 'unknown') + if isinstance(units, bytes): + units = units.decode() + + coordinates = None + if 'coordinates' in f: + coordinates = f['coordinates'][...] + + return TimeSeriesData( + time=time_vec, + data=data_matrix, + channels=channel_names, + sample_rate=sample_rate, + units=str(units), + coordinates=coordinates + ) + + +def _load_time_series_matlab(file_path: Path, **kwargs) -> TimeSeriesData: + """Load time series from MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + data = sio.loadmat(str(file_path)) + + time_vec = data.get('T', data.get('time', data.get('t', np.array([])))).flatten() + data_matrix = data.get('Y', data.get('data', data.get('Bdot', np.array([])))) + + if data_matrix.ndim == 1: + data_matrix = data_matrix.reshape(-1, 1) + + # Get channel names + channel_names = [] + if 'ynames' in data: + ynames = data['ynames'] + if ynames.dtype.names: # Structured array + channel_names = [str(ynames[i][0][0]) for i in range(len(ynames))] + else: + channel_names = [str(name[0]) if hasattr(name, '__len__') and len(name) > 0 else f"channel_{i+1}" + for i, name in enumerate(ynames.flatten())] + else: + channel_names = [f"channel_{i+1}" for i in range(data_matrix.shape[1])] + + sample_rate = float(1.0 / np.mean(np.diff(time_vec))) if len(time_vec) > 1 else 1.0 + + coordinates = None + if 'Yxy' in data: + coordinates = data['Yxy'] + + return TimeSeriesData( + time=time_vec, + data=data_matrix, + channels=channel_names, + sample_rate=sample_rate, + coordinates=coordinates + ) + + +def _load_config_json(file_path: Path) -> ConfigData: + """Load configuration from JSON file.""" + with open(file_path, 'r') as f: + data = json.load(f) + + return ConfigData( + analysis_params=data.get('analysis_params', {}), + processing_params=data.get('processing_params', {}), + plot_params=data.get('plot_params', {}), + io_params=data.get('io_params', {}), + file_paths=data.get('file_paths', {}), + user_settings=data.get('user_settings', {}), + version=data.get('version', '1.0') + ) + + +def _load_config_matlab(file_path: Path) -> ConfigData: + """Load configuration from MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + data = sio.loadmat(str(file_path)) + + # Convert MATLAB structure to dictionary + config_dict = {} + for key, value in data.items(): + if not key.startswith('__'): + if isinstance(value, np.ndarray) and value.dtype.names: + # Structured array (MATLAB struct) + config_dict[key] = _matlab_struct_to_dict(value) + else: + config_dict[key] = value + + return ConfigData( + analysis_params=config_dict.get('analysis_params', {}), + processing_params=config_dict.get('processing_params', {}), + plot_params=config_dict.get('plot_params', {}), + io_params=config_dict.get('io_params', {}), + file_paths=config_dict.get('file_paths', {}), + user_settings=config_dict.get('user_settings', {}), + version=str(config_dict.get('version', '1.0')) + ) + + +def _load_config_pickle(file_path: Path) -> ConfigData: + """Load configuration from pickle file.""" + with open(file_path, 'rb') as f: + data = pickle.load(f) + + if isinstance(data, ConfigData): + return data + elif isinstance(data, dict): + return ConfigData(**data) + else: + raise ValueError("Invalid pickle file format for ConfigData") + + +def _load_results_hdf5(file_path: Path) -> AnalysisResult: + """Load analysis results from HDF5 file.""" + if not HAS_HDF5: + raise ImportError("h5py is required for HDF5 support") + + with h5py.File(file_path, 'r') as f: + return AnalysisResult( + analysis_type=f.attrs.get('analysis_type', 'unknown').decode() if isinstance(f.attrs.get('analysis_type'), bytes) else str(f.attrs.get('analysis_type', 'unknown')), + shot_number=int(f.attrs.get('shot_number', 0)), + time_blocks=list(f['time_blocks'][...]), + frequencies=f['frequencies'][...], + eigenvalues=f['eigenvalues'][...], + mode_shapes=f['mode_shapes'][...], + stability=f['stability'][...], + parameters=dict(f.attrs) if f.attrs else {} + ) + + +def _load_results_matlab(file_path: Path) -> AnalysisResult: + """Load analysis results from MATLAB file.""" + if not HAS_SCIPY_IO: + raise ImportError("scipy is required for MATLAB file support") + + data = sio.loadmat(str(file_path)) + + return AnalysisResult( + analysis_type=str(data.get('analysis_type', ['unknown'])[0]), + shot_number=int(data.get('shot_number', [0])[0]), + time_blocks=data.get('time_blocks', []).tolist(), + frequencies=data.get('frequencies', np.array([])), + eigenvalues=data.get('eigenvalues', np.array([])), + mode_shapes=data.get('mode_shapes', np.array([])), + stability=data.get('stability', np.array([])), + parameters=_extract_matlab_params(data) + ) + + +def _load_results_pickle(file_path: Path) -> AnalysisResult: + """Load analysis results from pickle file.""" + with open(file_path, 'rb') as f: + data = pickle.load(f) + + if isinstance(data, AnalysisResult): + return data + else: + raise ValueError("Invalid pickle file format for AnalysisResult") + + +def _matlab_struct_to_dict(struct_array: np.ndarray) -> Dict[str, Any]: + """Convert MATLAB struct array to Python dictionary.""" + if struct_array.size == 0: + return {} + + result = {} + struct = struct_array.flat[0] + + for field_name in struct_array.dtype.names: + field_data = struct[field_name] + if isinstance(field_data, np.ndarray): + if field_data.dtype.names: # Nested struct + result[field_name] = _matlab_struct_to_dict(field_data) + else: + result[field_name] = field_data + else: + result[field_name] = field_data + + return result + + +def _extract_matlab_params(data: Dict[str, Any]) -> Dict[str, Any]: + """Extract parameters from MATLAB data structure.""" + params = {} + + # Look for common parameter fields + param_fields = ['parameters', 'params', 'options', 'config'] + + for field in param_fields: + if field in data: + field_data = data[field] + if isinstance(field_data, np.ndarray) and field_data.dtype.names: + params.update(_matlab_struct_to_dict(field_data)) + else: + params[field] = field_data + + return params \ No newline at end of file diff --git a/src/tokeye/eigspec/io/utils.py b/src/tokeye/eigspec/io/utils.py new file mode 100644 index 0000000..a20acc8 --- /dev/null +++ b/src/tokeye/eigspec/io/utils.py @@ -0,0 +1,781 @@ +""" +Data utilities for eigspec I/O operations. + +This module provides utility functions for data validation, manipulation, +merging, filtering, and bundling operations commonly used in plasma physics +data processing workflows. +""" + +from typing import Optional, Dict, List, Any, Union, Tuple +import numpy as np +import numpy.typing as npt +from pathlib import Path +import warnings + +from .data_structures import ( + ShotData, + ProcessedData, + AnalysisResult, + TimeSeriesData, + MirnovData, + ConfigData +) + + +def validate_data( + data: Union[TimeSeriesData, ProcessedData, ShotData], + strict: bool = False +) -> Tuple[bool, List[str]]: + """Validate data consistency and integrity. + + Args: + data: Data object to validate + strict: Whether to apply strict validation rules + + Returns: + Tuple of (is_valid, warning_messages) + + Raises: + ValueError: If critical validation errors are found + """ + warnings_list = [] + + if isinstance(data, TimeSeriesData): + warnings_list.extend(_validate_time_series(data, strict)) + elif isinstance(data, ProcessedData): + warnings_list.extend(_validate_processed_data(data, strict)) + elif isinstance(data, ShotData): + warnings_list.extend(_validate_shot_data(data, strict)) + else: + raise ValueError(f"Unsupported data type for validation: {type(data)}") + + # Check for critical errors vs warnings + critical_errors = [w for w in warnings_list if "ERROR:" in w] + + if critical_errors and strict: + raise ValueError(f"Critical validation errors: {critical_errors}") + + return len(critical_errors) == 0, warnings_list + + +def merge_datasets( + datasets: List[Union[TimeSeriesData, ProcessedData]], + method: str = "concatenate", + axis: str = "time" +) -> Union[TimeSeriesData, ProcessedData]: + """Merge multiple datasets along specified axis. + + Args: + datasets: List of datasets to merge + method: Merge method ('concatenate', 'average', 'stack') + axis: Axis along which to merge ('time', 'channels') + + Returns: + Merged dataset + + Raises: + ValueError: If datasets are incompatible for merging + """ + if not datasets: + raise ValueError("No datasets provided for merging") + + if len(datasets) == 1: + return datasets[0] + + # Check compatibility + _check_merge_compatibility(datasets, axis) + + if method == "concatenate": + return _concatenate_datasets(datasets, axis) + elif method == "average": + return _average_datasets(datasets) + elif method == "stack": + return _stack_datasets(datasets, axis) + else: + raise ValueError(f"Unknown merge method: {method}") + + +def filter_bad_channels( + data: Union[TimeSeriesData, ProcessedData, ShotData], + bad_channels: List[str], + in_place: bool = False +) -> Union[TimeSeriesData, ProcessedData, ShotData]: + """Remove bad channels from data. + + Args: + data: Data object to filter + bad_channels: List of channel names to remove + in_place: Whether to modify data in place + + Returns: + Filtered data object + """ + if not bad_channels: + return data if in_place else _copy_data_object(data) + + if isinstance(data, TimeSeriesData): + return _filter_time_series_channels(data, bad_channels, in_place) + elif isinstance(data, ProcessedData): + return _filter_processed_data_channels(data, bad_channels, in_place) + elif isinstance(data, ShotData): + return _filter_shot_data_channels(data, bad_channels, in_place) + else: + raise ValueError(f"Unsupported data type for channel filtering: {type(data)}") + + +def bundle_shot_data( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + coordinates: npt.NDArray[np.floating], + channel_names: List[str], + shot_number: int, + bad_channels: Optional[List[str]] = None, + interpolation_method: str = "linear", + detrend_method: str = "linear" +) -> ProcessedData: + """Bundle shot data into ProcessedData format. + + Similar to MATLAB bundle_shot_data function. + + Args: + time: Time vector + data: Data matrix (time x channels) + coordinates: Channel coordinates + channel_names: Channel names + shot_number: Shot number + bad_channels: Channels to exclude + interpolation_method: Method used for interpolation + detrend_method: Method used for detrending + + Returns: + ProcessedData object with bundled data + """ + if bad_channels is None: + bad_channels = [] + + # Validate input dimensions + if len(time) != data.shape[0]: + raise ValueError("Time vector length must match data rows") + if len(channel_names) != data.shape[1]: + raise ValueError("Number of channel names must match data columns") + if len(coordinates) != len(channel_names): + raise ValueError("Coordinates length must match number of channels") + + # Filter out bad channels + good_indices = [] + good_channels = [] + filtered_bad = [] + + for i, ch_name in enumerate(channel_names): + if ch_name in bad_channels: + print(f"*** removing: pointname \"{ch_name}\"; found at index #{i+1}.") + filtered_bad.append(ch_name) + else: + good_indices.append(i) + good_channels.append(ch_name) + + # Extract good channels + if good_indices: + filtered_data = data[:, good_indices] + filtered_coords = coordinates[good_indices] + else: + raise ValueError("No good channels remaining after filtering") + + return ProcessedData( + shot_number=shot_number, + time=time, + data=filtered_data, + coordinates=filtered_coords, + channel_names=good_channels, + bad_channels=filtered_bad, + processing_steps=["channel_filtering"], + interpolation_method=interpolation_method, + detrend_method=detrend_method + ) + + +def resample_time_series( + data: TimeSeriesData, + new_sample_rate: float, + method: str = "linear" +) -> TimeSeriesData: + """Resample time series data to new sampling rate. + + Args: + data: Input time series data + new_sample_rate: Target sampling rate (Hz) + method: Interpolation method ('linear', 'cubic', 'nearest') + + Returns: + Resampled time series data + """ + # Create new time vector + time_span = data.time[-1] - data.time[0] + num_new_samples = int(time_span * new_sample_rate) + 1 + new_time = np.linspace(data.time[0], data.time[-1], num_new_samples) + + # Interpolate data + new_data = np.zeros((len(new_time), data.data.shape[1])) + + for ch_idx in range(data.data.shape[1]): + new_data[:, ch_idx] = np.interp(new_time, data.time, data.data[:, ch_idx]) + + return TimeSeriesData( + time=new_time, + data=new_data, + channels=data.channels.copy(), + sample_rate=new_sample_rate, + units=data.units, + coordinates=data.coordinates.copy() if data.coordinates is not None else None, + metadata=data.metadata.copy() + ) + + +def align_time_series( + datasets: List[TimeSeriesData], + method: str = "intersection" +) -> List[TimeSeriesData]: + """Align multiple time series to common time base. + + Args: + datasets: List of time series to align + method: Alignment method ('intersection', 'union', 'first') + + Returns: + List of aligned time series + """ + if not datasets: + return [] + + if len(datasets) == 1: + return datasets + + # Determine common time base + if method == "intersection": + # Use intersection of all time ranges + start_time = max(data.time[0] for data in datasets) + end_time = min(data.time[-1] for data in datasets) + elif method == "union": + # Use union of all time ranges + start_time = min(data.time[0] for data in datasets) + end_time = max(data.time[-1] for data in datasets) + elif method == "first": + # Use time base of first dataset + start_time = datasets[0].time[0] + end_time = datasets[0].time[-1] + else: + raise ValueError(f"Unknown alignment method: {method}") + + # Find common sampling rate (use minimum for safety) + common_sample_rate = min(data.sample_rate for data in datasets) + + # Create common time vector + time_span = end_time - start_time + num_samples = int(time_span * common_sample_rate) + 1 + common_time = np.linspace(start_time, end_time, num_samples) + + # Interpolate all datasets to common time base + aligned_datasets = [] + for data in datasets: + new_data = np.zeros((len(common_time), data.data.shape[1])) + + for ch_idx in range(data.data.shape[1]): + new_data[:, ch_idx] = np.interp(common_time, data.time, data.data[:, ch_idx]) + + aligned_data = TimeSeriesData( + time=common_time, + data=new_data, + channels=data.channels.copy(), + sample_rate=common_sample_rate, + units=data.units, + coordinates=data.coordinates.copy() if data.coordinates is not None else None, + metadata=data.metadata.copy() + ) + aligned_datasets.append(aligned_data) + + return aligned_datasets + + +def compute_data_statistics( + data: Union[TimeSeriesData, ProcessedData] +) -> Dict[str, Any]: + """Compute statistical summary of data. + + Args: + data: Data object to analyze + + Returns: + Dictionary with statistical information + """ + stats = {} + + # Basic statistics + stats['mean'] = np.mean(data.data, axis=0) + stats['std'] = np.std(data.data, axis=0) + stats['min'] = np.min(data.data, axis=0) + stats['max'] = np.max(data.data, axis=0) + stats['median'] = np.median(data.data, axis=0) + + # Data quality metrics + stats['num_samples'] = data.data.shape[0] + stats['num_channels'] = data.data.shape[1] + stats['nan_count'] = np.sum(np.isnan(data.data), axis=0) + stats['inf_count'] = np.sum(np.isinf(data.data), axis=0) + + # Time statistics + if hasattr(data, 'time'): + stats['time_span'] = float(data.time[-1] - data.time[0]) + stats['sample_rate'] = float(1.0 / np.mean(np.diff(data.time))) + stats['time_gaps'] = _detect_time_gaps(data.time) + + # Signal quality metrics + stats['snr_estimate'] = _estimate_snr(data.data) + stats['dynamic_range'] = stats['max'] - stats['min'] + + return stats + + +def detect_outliers( + data: npt.NDArray[np.floating], + method: str = "iqr", + threshold: float = 3.0 +) -> npt.NDArray[np.bool_]: + """Detect outliers in data. + + Args: + data: Data array (samples x channels) + method: Detection method ('iqr', 'zscore', 'mad') + threshold: Threshold for outlier detection + + Returns: + Boolean array indicating outliers + """ + if method == "iqr": + return _detect_outliers_iqr(data, threshold) + elif method == "zscore": + return _detect_outliers_zscore(data, threshold) + elif method == "mad": + return _detect_outliers_mad(data, threshold) + else: + raise ValueError(f"Unknown outlier detection method: {method}") + + +def create_channel_map( + channel_names: List[str], + coordinates: Optional[npt.NDArray[np.floating]] = None +) -> Dict[str, Dict[str, Any]]: + """Create a mapping of channel information. + + Args: + channel_names: List of channel names + coordinates: Optional channel coordinates + + Returns: + Dictionary mapping channel names to information + """ + channel_map = {} + + for i, ch_name in enumerate(channel_names): + ch_info = { + 'index': i, + 'name': ch_name + } + + if coordinates is not None and i < len(coordinates): + ch_info['coordinates'] = { + 'R': float(coordinates[i, 0]) if coordinates.shape[1] > 0 else 0.0, + 'Z': float(coordinates[i, 1]) if coordinates.shape[1] > 1 else 0.0, + 'phi': float(coordinates[i, 2]) if coordinates.shape[1] > 2 else 0.0 + } + + channel_map[ch_name] = ch_info + + return channel_map + + +# Private helper functions +def _validate_time_series(data: TimeSeriesData, strict: bool) -> List[str]: + """Validate TimeSeriesData object.""" + warnings_list = [] + + # Check time vector + if len(data.time) == 0: + warnings_list.append("ERROR: Empty time vector") + elif len(np.unique(data.time)) != len(data.time): + warnings_list.append("WARNING: Non-unique time values detected") + elif not np.all(np.diff(data.time) > 0): + warnings_list.append("ERROR: Time vector is not monotonically increasing") + + # Check data matrix + if data.data.size == 0: + warnings_list.append("ERROR: Empty data matrix") + elif np.any(np.isnan(data.data)): + nan_count = np.sum(np.isnan(data.data)) + warnings_list.append(f"WARNING: {nan_count} NaN values in data") + elif np.any(np.isinf(data.data)): + inf_count = np.sum(np.isinf(data.data)) + warnings_list.append(f"WARNING: {inf_count} infinite values in data") + + # Check sampling rate consistency + if len(data.time) > 1: + actual_rate = 1.0 / np.mean(np.diff(data.time)) + rate_diff = abs(actual_rate - data.sample_rate) / data.sample_rate + if rate_diff > 0.05: # 5% tolerance + warnings_list.append(f"WARNING: Sample rate mismatch - stated: {data.sample_rate:.2f}, actual: {actual_rate:.2f}") + + return warnings_list + + +def _validate_processed_data(data: ProcessedData, strict: bool) -> List[str]: + """Validate ProcessedData object.""" + warnings_list = [] + + # Inherit time series validations + time_series_like = TimeSeriesData( + time=data.time, + data=data.data, + channels=data.channel_names, + sample_rate=1.0 / np.mean(np.diff(data.time)) if len(data.time) > 1 else 1.0 + ) + warnings_list.extend(_validate_time_series(time_series_like, strict)) + + # Additional processed data checks + if len(data.channel_names) != data.data.shape[1]: + warnings_list.append("ERROR: Channel names count mismatch with data columns") + + if len(data.coordinates) != len(data.channel_names): + warnings_list.append("ERROR: Coordinates count mismatch with channels") + + return warnings_list + + +def _validate_shot_data(data: ShotData, strict: bool) -> List[str]: + """Validate ShotData object.""" + warnings_list = [] + + # Check shot number + if data.shot_number <= 0: + warnings_list.append("WARNING: Invalid shot number") + + # Check time range + if data.time_range[0] >= data.time_range[1]: + warnings_list.append("ERROR: Invalid time range") + + # Check channel consistency + if len(data.channel_names) != len(data.channels): + warnings_list.append("ERROR: Channel names count mismatch with channel data") + + return warnings_list + + +def _check_merge_compatibility( + datasets: List[Union[TimeSeriesData, ProcessedData]], + axis: str +) -> None: + """Check if datasets are compatible for merging.""" + if axis == "time": + # Check that all datasets have same number of channels + ref_channels = len(datasets[0].channels if hasattr(datasets[0], 'channels') else datasets[0].channel_names) + for i, data in enumerate(datasets[1:], start=1): + data_channels = len(data.channels if hasattr(data, 'channels') else data.channel_names) + if data_channels != ref_channels: + raise ValueError(f"Dataset {i} has {data_channels} channels, expected {ref_channels}") + + elif axis == "channels": + # Check that all datasets have same time length + ref_time_len = len(datasets[0].time) + for i, data in enumerate(datasets[1:], start=1): + if len(data.time) != ref_time_len: + raise ValueError(f"Dataset {i} has {len(data.time)} time samples, expected {ref_time_len}") + + +def _concatenate_datasets( + datasets: List[Union[TimeSeriesData, ProcessedData]], + axis: str +) -> Union[TimeSeriesData, ProcessedData]: + """Concatenate datasets along specified axis.""" + if axis == "time": + # Concatenate along time axis + merged_time = np.concatenate([data.time for data in datasets]) + merged_data = np.concatenate([data.data for data in datasets], axis=0) + + # Use first dataset as template + template = datasets[0] + if isinstance(template, TimeSeriesData): + return TimeSeriesData( + time=merged_time, + data=merged_data, + channels=template.channels.copy(), + sample_rate=template.sample_rate, + units=template.units, + coordinates=template.coordinates.copy() if template.coordinates is not None else None, + metadata=template.metadata.copy() + ) + else: # ProcessedData + return ProcessedData( + shot_number=template.shot_number, + time=merged_time, + data=merged_data, + coordinates=template.coordinates.copy(), + channel_names=template.channel_names.copy(), + bad_channels=template.bad_channels.copy(), + processing_steps=template.processing_steps.copy() + ["time_concatenation"], + interpolation_method=template.interpolation_method, + detrend_method=template.detrend_method + ) + + elif axis == "channels": + # Concatenate along channels axis + merged_data = np.concatenate([data.data for data in datasets], axis=1) + + # Merge channel information + all_channels = [] + all_coords = [] + + for data in datasets: + if hasattr(data, 'channels'): + all_channels.extend(data.channels) + else: + all_channels.extend(data.channel_names) + + if hasattr(data, 'coordinates') and data.coordinates is not None: + all_coords.append(data.coordinates) + + merged_coords = np.concatenate(all_coords) if all_coords else None + + # Use first dataset as template + template = datasets[0] + if isinstance(template, TimeSeriesData): + return TimeSeriesData( + time=template.time.copy(), + data=merged_data, + channels=all_channels, + sample_rate=template.sample_rate, + units=template.units, + coordinates=merged_coords, + metadata=template.metadata.copy() + ) + else: # ProcessedData + return ProcessedData( + shot_number=template.shot_number, + time=template.time.copy(), + data=merged_data, + coordinates=merged_coords, + channel_names=all_channels, + bad_channels=template.bad_channels.copy(), + processing_steps=template.processing_steps.copy() + ["channel_concatenation"], + interpolation_method=template.interpolation_method, + detrend_method=template.detrend_method + ) + else: + raise ValueError(f"Unknown concatenation axis: {axis}") + + +def _average_datasets( + datasets: List[Union[TimeSeriesData, ProcessedData]] +) -> Union[TimeSeriesData, ProcessedData]: + """Average datasets (requires same time base).""" + # First align all datasets to same time base + aligned_datasets = align_time_series(datasets) if all(isinstance(d, TimeSeriesData) for d in datasets) else datasets + + # Average data + data_arrays = [data.data for data in aligned_datasets] + averaged_data = np.mean(data_arrays, axis=0) + + # Use first dataset as template + template = aligned_datasets[0] + if isinstance(template, TimeSeriesData): + return TimeSeriesData( + time=template.time.copy(), + data=averaged_data, + channels=template.channels.copy(), + sample_rate=template.sample_rate, + units=template.units, + coordinates=template.coordinates.copy() if template.coordinates is not None else None, + metadata=template.metadata.copy() + ) + else: # ProcessedData + return ProcessedData( + shot_number=template.shot_number, + time=template.time.copy(), + data=averaged_data, + coordinates=template.coordinates.copy(), + channel_names=template.channel_names.copy(), + bad_channels=template.bad_channels.copy(), + processing_steps=template.processing_steps.copy() + ["ensemble_averaging"], + interpolation_method=template.interpolation_method, + detrend_method=template.detrend_method + ) + + +def _stack_datasets( + datasets: List[Union[TimeSeriesData, ProcessedData]], + axis: str +) -> Union[TimeSeriesData, ProcessedData]: + """Stack datasets (adds new dimension).""" + # This is a placeholder - actual implementation would depend on specific requirements + return _concatenate_datasets(datasets, axis) + + +def _copy_data_object( + data: Union[TimeSeriesData, ProcessedData, ShotData] +) -> Union[TimeSeriesData, ProcessedData, ShotData]: + """Create a copy of data object.""" + if isinstance(data, TimeSeriesData): + return TimeSeriesData( + time=data.time.copy(), + data=data.data.copy(), + channels=data.channels.copy(), + sample_rate=data.sample_rate, + units=data.units, + coordinates=data.coordinates.copy() if data.coordinates is not None else None, + metadata=data.metadata.copy() + ) + elif isinstance(data, ProcessedData): + return ProcessedData( + shot_number=data.shot_number, + time=data.time.copy(), + data=data.data.copy(), + coordinates=data.coordinates.copy(), + channel_names=data.channel_names.copy(), + bad_channels=data.bad_channels.copy(), + processing_steps=data.processing_steps.copy(), + interpolation_method=data.interpolation_method, + detrend_method=data.detrend_method + ) + else: # ShotData + return ShotData( + shot_number=data.shot_number, + time_range=data.time_range, + channels={k: v.copy() for k, v in data.channels.items()}, + coordinates=data.coordinates.copy(), + channel_names=data.channel_names.copy(), + acquisition_info=data.acquisition_info.copy(), + preprocessing=data.preprocessing.copy(), + quality_flags=data.quality_flags.copy() + ) + + +def _filter_time_series_channels( + data: TimeSeriesData, + bad_channels: List[str], + in_place: bool +) -> TimeSeriesData: + """Filter channels from TimeSeriesData.""" + if not in_place: + data = _copy_data_object(data) + + # Find indices of good channels + good_indices = [i for i, ch in enumerate(data.channels) if ch not in bad_channels] + + if not good_indices: + raise ValueError("No good channels remaining after filtering") + + # Filter data + data.data = data.data[:, good_indices] + data.channels = [data.channels[i] for i in good_indices] + + if data.coordinates is not None: + data.coordinates = data.coordinates[good_indices] + + return data + + +def _filter_processed_data_channels( + data: ProcessedData, + bad_channels: List[str], + in_place: bool +) -> ProcessedData: + """Filter channels from ProcessedData.""" + if not in_place: + data = _copy_data_object(data) + + # Find indices of good channels + good_indices = [i for i, ch in enumerate(data.channel_names) if ch not in bad_channels] + + if not good_indices: + raise ValueError("No good channels remaining after filtering") + + # Filter data + data.data = data.data[:, good_indices] + data.coordinates = data.coordinates[good_indices] + data.channel_names = [data.channel_names[i] for i in good_indices] + + # Update bad channels list + data.bad_channels.extend([ch for ch in bad_channels if ch not in data.bad_channels]) + data.processing_steps.append("channel_filtering") + + return data + + +def _filter_shot_data_channels( + data: ShotData, + bad_channels: List[str], + in_place: bool +) -> ShotData: + """Filter channels from ShotData.""" + if not in_place: + data = _copy_data_object(data) + + # Remove bad channels + for ch in bad_channels: + if ch in data.channels: + del data.channels[ch] + if ch in data.coordinates: + del data.coordinates[ch] + if ch in data.channel_names: + data.channel_names.remove(ch) + + return data + + +def _detect_time_gaps(time: npt.NDArray[np.floating], threshold: float = 2.0) -> List[int]: + """Detect gaps in time series.""" + dt = np.diff(time) + median_dt = np.median(dt) + gap_indices = np.where(dt > threshold * median_dt)[0] + return gap_indices.tolist() + + +def _estimate_snr(data: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: + """Estimate signal-to-noise ratio for each channel.""" + # Simple SNR estimate using signal variance vs noise estimate + signal_power = np.var(data, axis=0) + + # Estimate noise from high-frequency content (crude approximation) + if data.shape[0] > 10: + noise_estimate = np.var(np.diff(data, axis=0), axis=0) / 2.0 + snr = signal_power / (noise_estimate + 1e-12) # Avoid division by zero + else: + snr = np.ones(data.shape[1]) # Fallback + + return snr + + +def _detect_outliers_iqr(data: npt.NDArray[np.floating], threshold: float) -> npt.NDArray[np.bool_]: + """Detect outliers using IQR method.""" + q1 = np.percentile(data, 25, axis=0) + q3 = np.percentile(data, 75, axis=0) + iqr = q3 - q1 + + lower_bound = q1 - threshold * iqr + upper_bound = q3 + threshold * iqr + + outliers = (data < lower_bound) | (data > upper_bound) + return outliers + + +def _detect_outliers_zscore(data: npt.NDArray[np.floating], threshold: float) -> npt.NDArray[np.bool_]: + """Detect outliers using Z-score method.""" + z_scores = np.abs((data - np.mean(data, axis=0)) / np.std(data, axis=0)) + return z_scores > threshold + + +def _detect_outliers_mad(data: npt.NDArray[np.floating], threshold: float) -> npt.NDArray[np.bool_]: + """Detect outliers using Median Absolute Deviation method.""" + median = np.median(data, axis=0) + mad = np.median(np.abs(data - median), axis=0) + + # Modified Z-score using MAD + modified_z_scores = 0.6745 * (data - median) / (mad + 1e-12) + return np.abs(modified_z_scores) > threshold \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/__init__.py b/src/tokeye/eigspec/utils/__init__.py new file mode 100644 index 0000000..578df1d --- /dev/null +++ b/src/tokeye/eigspec/utils/__init__.py @@ -0,0 +1,3 @@ +""" +Core functionality for eigspec package. +""" \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/autoregressive_pca.py b/src/tokeye/eigspec/utils/autoregressive_pca.py new file mode 100644 index 0000000..d322eb5 --- /dev/null +++ b/src/tokeye/eigspec/utils/autoregressive_pca.py @@ -0,0 +1,412 @@ +""" +Autoregressive Principal Component Analysis (AR-PCA) for eigspec package. + +This module provides AR-PCA algorithms for time series modeling and spectral analysis: +- Autoregressive modeling with PCA dimensionality reduction +- Eigenvalue-based frequency and damping estimation +- Model selection and validation procedures + +Based on the MATLAB eigspec toolbox AR-PCA functions: +- arpca.m - Main AR-PCA algorithm implementation +- ssi1cax.m - Extended SSI algorithm with AR-PCA components +- Various supporting utilities for model fitting and validation +""" + +from dataclasses import dataclass +from typing import List, Optional, Union + +import numpy as np +from numpy.typing import NDArray + + +@dataclass +class ARPCAResult: + """ + Container for AR/PCA analysis results. + + Attributes: + A: State transition matrix, shape (n_states, n_states) + K: Kalman gain matrix, shape (n_states, n_outputs) + C: Output matrix, shape (n_outputs, n_states) + H: AR model Markov blocks, shape (n_outputs, n_outputs * past) + pca: PCA eigenvalues (if dimension reduction was used) + models: List of models for multiple orders (if applicable) + Ry: Data covariance matrix (if residual computed) + Re: Residual error covariance (if residual computed, single model case) + """ + A: Optional[NDArray[np.floating]] + K: Optional[NDArray[np.floating]] + C: Optional[NDArray[np.floating]] + H: NDArray[np.floating] + pca: Optional[NDArray[np.floating]] + models: Optional[List['ARPCAModel']] + Ry: Optional[NDArray[np.floating]] = None + Re: Optional[NDArray[np.floating]] = None + + +@dataclass +class ARPCAModel: + """ + Container for individual AR/PCA model. + + Attributes: + A: State transition matrix + K: Kalman gain matrix + C: Output matrix + Re: Residual error covariance (optional) + """ + A: NDArray[np.floating] + K: NDArray[np.floating] + C: NDArray[np.floating] + Re: Optional[NDArray[np.floating]] = None + + +def arpca( + y: NDArray[np.number], + p: int, + r: Union[int, float, List[Union[int, float]]], + compute_residual: bool = False +) -> ARPCAResult: + """ + AR/PCA time-series modeling for a block of multichannel data. + + This function implements AR/PCA modeling where the past horizon for the linear + predictor is p samples. The state order r can be specified in several ways: + - Scalar 1 <= r <= m*p: Return single model with order r + - Vector r: Return list of models with orders in r + - Scalar 0 < r < 1: Select order based on PCA eigenvalue energy fraction + - Scalar r <= 0: Return unreduced lag-p AR model + + Args: + y: Multichannel data, shape (n_samples, n_channels) + p: Past horizon for linear predictor (positive integer) + r: State order specification (see above) + compute_residual: If True, compute residual error covariance + + Returns: + ARPCAResult containing system matrices and analysis results + + Raises: + ValueError: If parameters are invalid or data dimensions are problematic + + Example: + >>> # Generate AR(2) process + >>> np.random.seed(42) + >>> y = np.random.randn(1000, 3) # 3 channels, 1000 samples + >>> result = arpca(y, p=10, r=5) # Past=10, reduced order=5 + >>> print(f"A matrix shape: {result.A.shape}") + >>> print(f"Kalman gain shape: {result.K.shape}") + """ + if not isinstance(y, np.ndarray): + raise TypeError("y must be a numpy array") + if y.ndim != 2: + raise ValueError(f"y must be 2D, got shape {y.shape}") + if not isinstance(p, int) or p < 1: + raise ValueError("p must be a positive integer") + + N, ny = y.shape + + if ny >= N: + raise ValueError(f"Number of channels {ny} exceeds number of samples {N}") + + # Ensure we have enough samples for the lagged structure + # Need at least p+1 samples to form one prediction instance + if p >= N: + raise ValueError(f"Lag p={p} must be less than batch length {N}") + + # Warn if lag is very large relative to data length + if p >= N // 2: + import warnings + warnings.warn(f"Large lag p={p} relative to batch length {N} may lead to poor estimates") + + # Main AR/PCA computation + result = _subarpca(y, p, r) + + # Compute residual error covariance if requested + if compute_residual: + result = _compute_residual_error(result, y, p) + + return result + + +def _subarpca( + signal: NDArray[np.number], + past_horizon: int, + order: Union[int, float, List[Union[int, float]]] +) -> ARPCAResult: + """ + Core AR/PCA computation. + + Args: + signal: Input data, shape (n_samples, n_channels) + past_horizon: Past horizon + order: Order specification + + Returns: + ARPCAResult with computed models + """ + n_samples, n_channels = signal.shape + n_prediction_instances = n_samples - past_horizon + n_lagged_dimension = n_channels * past_horizon + + # Build lagged data matrices + Z = signal.T # Transpose for easier indexing + Z_past = np.zeros((n_lagged_dimension, n_prediction_instances)) # Past data matrix + Y_current = np.zeros((n_channels, n_prediction_instances)) # Current output matrix + + for i in range(past_horizon, n_samples): # MATLAB: for kk=(p+1):N + # Build past vector: [y(k-1), y(k-2), ..., y(k-p)] + # Fixed indexing: collect past_horizon samples in reverse order + past_indices = list(range(i-1, i-1-past_horizon, -1)) # [i-1, i-2, ..., i-p] + past_block = Z[:, past_indices] # Shape: (n_channels, past_horizon) + Z_past[:, i-past_horizon] = past_block.flatten() + Y_current[:, i-past_horizon] = signal[i, :] + + # Compute lagged covariance and AR model + R_z_past = Z_past @ Z_past.T + H = (Y_current @ Z_past.T) @ np.linalg.pinv(R_z_past) # AR model, Markov blocks + + # Build full lag-p state-space model + AK, K, C = _build_lag_model(H, n_channels) + + if isinstance(order, (list, tuple, np.ndarray)): + # Multiple models requested + order_vec = list(order) + if any(ri <= 0 or ri > n_lagged_dimension for ri in order_vec if ri >= 1): + raise ValueError("Order elements must be in range (0,1) or [1,m*p]") + + V, D = _sorted_eig(R_z_past) + cumd = np.cumsum(D) / np.sum(D) # Cumulative energy + + models = [] + for ri in order_vec: + if 0 < ri < 1: + # Energy-based selection + ri = int(np.argmax(cumd >= ri) + 1) + + ri = int(ri) + V_r = V[:, :ri] + AKr, Kr, Cr = _deflate_model(H, V_r) + models.append(ARPCAModel(A=AKr + Kr @ Cr, K=Kr, C=Cr)) + + return ARPCAResult(A=None, K=None, C=None, H=H, pca=D, models=models) + + elif isinstance(order, (int, float)): + if order <= 0: + # Unreduced model + return ARPCAResult(A=AK + K @ C, K=K, C=C, H=H, pca=None, models=None) + + elif order >= 1 and order <= n_lagged_dimension: + # Single reduced model with specific order + order = int(order) + V, D = _sorted_eig(R_z_past) + V_r = V[:, :order] + AKr, Kr, Cr = _deflate_model(H, V_r) + return ARPCAResult(A=AKr + Kr @ Cr, K=Kr, C=Cr, H=H, pca=D, models=None) + + elif 0 < order < 1: + # Single reduced model based on energy fraction + V, D = _sorted_eig(R_z_past) + cumd = np.cumsum(D) / np.sum(D) + order = int(np.argmax(cumd >= order) + 1) + V_r = V[:, :order] + AKr, Kr, Cr = _deflate_model(H, V_r) + return ARPCAResult(A=AKr + Kr @ Cr, K=Kr, C=Cr, H=H, pca=D, models=None) + + else: + raise ValueError(f"Invalid order specification order={order}") + + else: + raise ValueError("order must be int, float, or list/array") + + +def _deflate_model( + H: NDArray[np.floating], + V_r: NDArray[np.floating] +) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: + """ + Compute reduced-order model matrices from full AR model. + + Args: + H: AR model Markov blocks, shape (m, m*p) + V_r: Reduction matrix, shape (m*p, r) + + Returns: + Tuple of (AKr, Kr, Cr) - reduced model matrices + """ + m = H.shape[0] + p = H.shape[1] // m + + # Extract matrices + Kr = V_r[:m, :].T # First m rows, transposed + Cr = H @ V_r # Apply reduction to Markov blocks + + # Compute reduced A matrix - Fixed to match MATLAB indexing + AKr = V_r[m:p*m, :].T @ V_r[:m*(p-1), :] + + return AKr, Kr, Cr + + +def _build_lag_model( + H: NDArray[np.floating], + nz: int +) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: + """ + Build full lag-p state-space model from AR Markov blocks. + + Args: + H: AR model Markov blocks, shape (nz, nz*p) + nz: Number of outputs/channels + + Returns: + Tuple of (A, B, C) where B is Kalman gain structure + """ + p = H.shape[1] // nz + + # Build companion form state matrix + A = np.block([ + [np.zeros((nz, p * nz))], + [np.block([np.eye((p-1) * nz), np.zeros(((p-1) * nz, nz))])] + ]) + + # Input matrix (Kalman gain structure) + B = np.block([ + [np.eye(nz)], + [np.zeros(((p-1) * nz, nz))] + ]) + + # Output matrix + C = H + + return A, B, C + + +def _sorted_eig(A: NDArray[np.floating]) -> tuple[NDArray[np.floating], NDArray[np.floating]]: + """ + Compute eigendecomposition with eigenvalues sorted in descending order. + + Args: + A: Matrix (may not be symmetric) + + Returns: + Tuple of (eigenvectors, eigenvalues) sorted by eigenvalue magnitude + """ + # Use general eigenvalue decomposition to match MATLAB eig() behavior + eigvals, eigvecs = np.linalg.eig(A) + + # Sort in descending order by real part (matching MATLAB behavior) + idx = np.argsort(np.real(eigvals))[::-1] + eigvals = eigvals[idx] + eigvecs = eigvecs[:, idx] + + return eigvecs, eigvals + + +def _compute_residual_error(result: ARPCAResult, y: NDArray[np.number], p: int) -> ARPCAResult: + """ + Compute residual error covariance for AR/PCA models. + + Args: + result: AR/PCA result containing models + y: Original data matrix, shape (n_samples, n_channels) + p: Past horizon + + Returns: + Updated ARPCAResult with residual error covariances + """ + from scipy.signal import lsim + + N, ny = y.shape + + # Compute data covariance + Ry = (y.T @ y) / N + + if hasattr(result, 'A') and result.A is not None: + # Single model case + A, K, C = result.A, result.K, result.C + + # Create discrete-time state-space model for residual computation + # H = (sI - (A - KC))^{-1} K with D = I for residual filter + # In discrete time: H(z) = C(zI - (A - KC))^{-1}K + I + try: + # Simulate the system to get residuals + # This is a simplified implementation - full MATLAB version uses ss() and lsim() + residuals = _compute_model_residuals(y, A, K, C, p) + Re = (residuals.T @ residuals) / (N - p) + + # Add residual info to result + result_dict = result.__dict__.copy() + result_dict['Ry'] = Ry + result_dict['Re'] = Re + result = ARPCAResult(**result_dict) + + except Exception: + # Fallback: skip residual computation if simulation fails + pass + + elif hasattr(result, 'models') and result.models: + # Multiple models case + for i, model in enumerate(result.models): + try: + residuals = _compute_model_residuals(y, model.A, model.K, model.C, p) + Re = (residuals.T @ residuals) / (N - p) + + # Add residual info to model (create new model with residual) + model_dict = model.__dict__.copy() + model_dict['Re'] = Re + result.models[i] = ARPCAModel(**model_dict) + + except Exception: + # Skip residual for this model if computation fails + continue + + # Add data covariance to main result + result_dict = result.__dict__.copy() + result_dict['Ry'] = Ry + result = ARPCAResult(**result_dict) + + return result + + +def _compute_model_residuals( + y: NDArray[np.number], + A: NDArray[np.floating], + K: NDArray[np.floating], + C: NDArray[np.floating], + p: int +) -> NDArray[np.floating]: + """ + Compute residuals for a single AR/PCA model. + + Args: + y: Data matrix, shape (n_samples, n_channels) + A: State transition matrix + K: Kalman gain matrix + C: Output matrix + p: Past horizon + + Returns: + Residuals matrix, shape (n_samples-p, n_channels) + """ + N, ny = y.shape + + # Simple residual computation using one-step-ahead prediction + # This is a simplified version of the MATLAB lsim approach + residuals = np.zeros((N - p, ny)) + + # Initialize state + n_states = A.shape[0] + x = np.zeros(n_states) + + for i in range(p, N): + # Predict output + y_pred = C @ x + + # Compute residual + residuals[i - p, :] = y[i, :] - y_pred + + # Update state: x[k+1] = A*x[k] + K*(y[k] - C*x[k]) + innovation = y[i, :] - y_pred + x = A @ x + K @ innovation + + return residuals \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/block_processing.py b/src/tokeye/eigspec/utils/block_processing.py new file mode 100644 index 0000000..2627cfe --- /dev/null +++ b/src/tokeye/eigspec/utils/block_processing.py @@ -0,0 +1,232 @@ +""" +Block-based data processing for eigspec package. + +This module provides block-based analysis functions for large dataset processing: +- Random projection block analysis with modal identification +- Time-windowed spectral analysis and feature extraction +- Block-based system identification workflows + +Based on the MATLAB eigspec toolbox block processing functions: +- rndspecx.m - Random projection spectral analysis for individual blocks +- eigspec_mmain.m - Main block-wise analysis workflow +- collect_rep_data.m - Feature collection across time blocks +- view_pcaspec_results.m - Block-based results visualization +""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np +from numpy.typing import NDArray + +from .modal_analysis import ModalShortlist, ShapeEstimates, order_mac +from .subspace_identification import covariance_driven_ssi, canonical_correlation_ssi +from .autoregressive_pca import arpca + +@dataclass +class BlockAnalysisResult: + """ + Container for block processing results. + + Attributes: + modal_analysis: Modal analysis results after dimension reduction + shape_estimates: Shape estimates after dimension reduction + processing_time: Time taken to process block in seconds + demean_applied: Whether block was demeaned before processing + time_step: Time step between samples in seconds + time_indices: (start_idx, end_idx) time indices of block + center_time: Center time of block in seconds + filter_time: Filter time in seconds + """ + modal_analysis: ModalShortlist + shape_estimates: Optional[ShapeEstimates] + processing_time: float + demean_applied: bool + time_step: float + time_indices: Tuple[int, int] + center_time: float + filter_time: float + +@dataclass +class RandomProjectionResult: + """ + Container for random projection block results. + + Attributes: + projection_matrix: Random projection matrix, shape (reduced_dim, n_channels) or None + modal_analysis: Modal analysis results after projection + shape_estimates: Shape estimates after projection + """ + projection_matrix: Optional[NDArray[np.floating]] + modal_analysis: ModalShortlist + shape_estimates: Optional[ShapeEstimates] + +def random_projection_block_analysis( + data_block: NDArray[np.number], + analysis_parameters: List[int], + matching_thresholds: Tuple[float, float], + use_canonical_correlation: bool = False, + random_seed: Optional[int] = None +) -> RandomProjectionResult: + """ + Process a data block using random projection and modal analysis. + + Args: + data_block: Input data block, shape (n_samples, n_channels) + analysis_parameters: Analysis configuration: + For SSI: [reduced_dim, future_horizon, past_horizon, order1, order2] + For AR/PCA: [reduced_dim, past_horizon, order1, order2] + where: + - reduced_dim: Target dimension (<0 for orthonormal, >0 for random, 0 for no projection) + - future_horizon/past_horizon: Number of future/past samples for SSI + - order1/order2: Model orders + matching_thresholds: (MAC threshold, distance threshold) for mode matching + use_canonical_correlation: Whether to use CCA/CVA in SSI algorithm + random_seed: Optional seed for random number generation + + Returns: + RandomProjectionResult containing: + - Random projection matrix (if used) + - Modal analysis results + - Shape estimates + + Raises: + ValueError: If reduced dimension is larger than number of channels + + Example: + >>> data = np.random.randn(1000, 10) # 1000 samples, 10 channels + >>> params = [-5, 10, 10, 2, 2] # SSI with orthonormal projection to 5D + >>> thresholds = (0.9, 0.9) # MAC and distance thresholds + >>> result = random_projection_block_analysis(data, params, thresholds) + """ + if data_block.ndim != 2: + raise ValueError(f"data_block must be 2D, got shape {data_block.shape}") + if not isinstance(analysis_parameters, list) or not all(isinstance(x, int) for x in analysis_parameters): + raise TypeError("analysis_parameters must be a list of integers") + if len(analysis_parameters) not in (4, 5): + raise ValueError("analysis_parameters must have 4 elements (AR/PCA) or 5 elements (SSI)") + + n_samples, n_channels = data_block.shape + reduced_dimension = analysis_parameters[0] + + if abs(reduced_dimension) > n_channels: + raise ValueError(f"Reduced dimension {abs(reduced_dimension)} cannot exceed input dimension {n_channels}") + + # Set random seed if provided + if random_seed is not None: + np.random.seed(random_seed) + + # TODO: Matrix with Random Variance. Orthonormal Projection with QR Factorization. Or Simpler way. 2 ways to do. + # See if returns Q or R matrix. Want Q matrix. + # Generate random projection matrix + if reduced_dimension < 0: + # Orthonormal projection + projection_matrix = np.linalg.qr(np.random.randn(n_channels, -reduced_dimension))[0].T + elif reduced_dimension > 0: + # Random projection - matching MATLAB behavior (no normalization) + projection_matrix = np.random.randn(reduced_dimension, n_channels) + else: + projection_matrix = None + + reduced_dimension = abs(reduced_dimension) + + # Apply projection if needed + projected_data = (projection_matrix @ data_block.T).T if projection_matrix is not None else data_block + + # Make this its own section + if len(analysis_parameters) == 4: + # AR/PCA mode: [reduced_dim, past_horizon, order1, order2] + past_horizon = analysis_parameters[1] + model_orders = analysis_parameters[2:4] + + # Call ARPCA algorithm + arpca_result = arpca(projected_data, past_horizon, list(model_orders)) + + if arpca_result.models: + # Multiple models case + first_model = arpca_result.models[0] + second_model = arpca_result.models[1] if len(arpca_result.models) > 1 else first_model + + modal_analysis = order_mac( + first_model.C, first_model.A, + second_model.C, second_model.A, + matching_thresholds + ) + else: + # Single model case - create dummy second model for OMAC comparison + state_matrix, output_matrix = arpca_result.A, arpca_result.C + if state_matrix is not None and output_matrix is not None: + # Small perturbation for comparison + perturbed_state = state_matrix * 1.01 + perturbed_output = output_matrix * 1.01 + else: + raise ValueError("ARPCA returned None matrices") + modal_analysis = order_mac( + output_matrix, state_matrix, + perturbed_output, perturbed_state, + matching_thresholds + ) + + else: + # SSI mode: [reduced_dim, future_horizon, past_horizon, order1, order2] + future_horizon = analysis_parameters[1] + past_horizon = analysis_parameters[2] + model_orders = analysis_parameters[3:5] + identification_params = [future_horizon, past_horizon] + model_orders + + if reduced_dimension == 0: + # Full-signal analysis + if use_canonical_correlation: + # More computationally heavy + ssi_result = canonical_correlation_ssi(projected_data, identification_params) + else: + # Use this by default + ssi_result = covariance_driven_ssi(projected_data, identification_params) + else: + # Reduced-signal analysis (same as full signal since data is already projected) + if use_canonical_correlation: + ssi_result = canonical_correlation_ssi(projected_data, identification_params) + else: + ssi_result = covariance_driven_ssi(projected_data, identification_params) + + if ssi_result.models and len(ssi_result.models) >= 2: + # Multiple models case + first_model = ssi_result.models[0] + second_model = ssi_result.models[1] + modal_analysis = order_mac( + first_model.output_matrix, first_model.state_matrix, + second_model.output_matrix, second_model.state_matrix, + matching_thresholds + ) + elif ssi_result.models and len(ssi_result.models) == 1: + # Single model case - create dummy second model + first_model = ssi_result.models[0] + perturbed_state = first_model.state_matrix * 1.01 + perturbed_output = first_model.output_matrix * 1.01 + modal_analysis = order_mac( + first_model.output_matrix, first_model.state_matrix, + perturbed_output, perturbed_state, + matching_thresholds + ) + else: + # Use single model results directly + state_matrix, output_matrix = ssi_result.state_matrix, ssi_result.output_matrix + if state_matrix is not None and output_matrix is not None: + perturbed_state, perturbed_output = state_matrix * 1.01, output_matrix * 1.01 + modal_analysis = order_mac( + output_matrix, state_matrix, + perturbed_output, perturbed_state, + matching_thresholds + ) + else: + raise ValueError("SSI returned None matrices") + + # Estimate shape vectors from original (unprojected) data + from .modal_analysis import shapes_from_freq + shape_estimates = shapes_from_freq(data_block, modal_analysis) + + return RandomProjectionResult( + projection_matrix=projection_matrix, + modal_analysis=modal_analysis, + shape_estimates=shape_estimates + ) diff --git a/src/tokeye/eigspec/utils/clustering.py b/src/tokeye/eigspec/utils/clustering.py new file mode 100644 index 0000000..eb7be11 --- /dev/null +++ b/src/tokeye/eigspec/utils/clustering.py @@ -0,0 +1,862 @@ +""" +Clustering analysis module for eigspec. + +This module provides clustering algorithms and pattern recognition tools specifically +designed for modal analysis applications, including: +- Distance/similarity metrics (Euclidean, Cosine, MAC) +- Clustering algorithms (K-means, Spectral clustering, Medoids) +- Graph construction methods (Full, kNN, mutual kNN, epsilon neighborhoods) +- Specialized modal analysis clustering utilities + +Based on the MATLAB eigspec toolbox clustering functions. +""" + +from typing import Union, Optional, Literal, cast, Dict +import numpy as np +import numpy.typing as npt +from scipy.linalg import eigh +from scipy.spatial.distance import pdist, squareform +from dataclasses import dataclass +import warnings + + +@dataclass +class ClusteringResult: + """Result from clustering analysis. + + Attributes: + labels: Cluster labels for each data point + centroids: Cluster centroids (for k-means) + medoid_indices: Indices of cluster medoids + cost: Final clustering cost/objective function value + eigenvalues: Eigenvalues from spectral clustering (if applicable) + iterations: Number of iterations to convergence + """ + labels: npt.NDArray[np.int32] + centroids: Optional[npt.NDArray[np.floating]] = None + medoid_indices: Optional[npt.NDArray[np.int32]] = None + cost: Optional[float] = None + eigenvalues: Optional[npt.NDArray[np.floating]] = None + iterations: Optional[int] = None + + +class DistanceMetric: + """Distance and similarity metrics for clustering analysis.""" + + @staticmethod + def euclidean_distance(x: npt.NDArray, y: npt.NDArray) -> float: + """Euclidean distance between two vectors.""" + return float(np.linalg.norm(x - y)) + + @staticmethod + def cosine_distance(x: npt.NDArray, y: npt.NDArray) -> float: + """Cosine distance between two real vectors.""" + if np.iscomplexobj(x) or np.iscomplexobj(y): + raise ValueError("Cosine distance not compatible with complex data") + + x_flat = x.ravel() + y_flat = y.ravel() + + norm_x = np.linalg.norm(x_flat) + norm_y = np.linalg.norm(y_flat) + + if norm_x == 0 or norm_y == 0: + return 1.0 + + cosine_sim = np.dot(x_flat, y_flat) / (norm_x * norm_y) + return float(1.0 - cosine_sim) + + @staticmethod + def mac_distance(x: npt.NDArray, y: npt.NDArray) -> float: + """Modal Assurance Criterion (MAC) distance.""" + numerator = np.abs(np.vdot(x, y)) ** 2 + denominator = (np.vdot(x, x) * np.vdot(y, y)) + + if denominator == 0: + return 1.0 + + mac_value = np.real(numerator / denominator) + return float(1.0 - mac_value) + + @staticmethod + def euclidean_similarity(x: npt.NDArray, y: npt.NDArray, sigma: float) -> float: + """Gaussian similarity based on Euclidean distance.""" + dist_sq = np.linalg.norm(x - y) ** 2 + return float(np.exp(-dist_sq / (2 * sigma**2))) + + @staticmethod + def cosine_similarity(x: npt.NDArray, y: npt.NDArray) -> float: + """Cosine similarity between two real vectors.""" + if np.iscomplexobj(x) or np.iscomplexobj(y): + raise ValueError("Cosine similarity not compatible with complex data") + + x_flat = x.ravel() + y_flat = y.ravel() + + norm_x = np.linalg.norm(x_flat) + norm_y = np.linalg.norm(y_flat) + + if norm_x == 0 or norm_y == 0: + return 0.0 + + return float(np.dot(x_flat, y_flat) / (norm_x * norm_y)) + + @staticmethod + def mac_similarity(x: npt.NDArray, y: npt.NDArray) -> float: + """Modal Assurance Criterion (MAC) similarity.""" + numerator = np.abs(np.vdot(x, y)) ** 2 + denominator = (np.vdot(x, x) * np.vdot(y, y)) + + if denominator == 0: + return 0.0 + + return float(np.real(numerator / denominator)) + + +def distance_matrix( + X: npt.NDArray[np.floating], + metric: Literal["euclidean", "cosine", "mac"] = "euclidean" +) -> npt.NDArray[np.floating]: + """Compute symmetric distance matrix between data points. + + Args: + X: Data matrix where each column is a feature vector + metric: Distance metric to use + + Returns: + Symmetric distance matrix + """ + n = X.shape[1] + + if metric == "euclidean": + # Use scipy's optimized pdist for Euclidean distance + distances = pdist(X.T, metric='euclidean') + return cast(npt.NDArray[np.floating], squareform(distances)) + + elif metric == "mac": + # Vectorized MAC distance computation + # Normalize columns for MAC computation + norms = np.sqrt(np.sum(X.conj() * X, axis=0, keepdims=True)) + norms[norms == 0] = 1 # Avoid division by zero + X_norm = X / norms + + # MAC similarity matrix + similarity = np.abs(X_norm.conj().T @ X_norm) ** 2 + + # Convert to distance and ensure zero diagonal + distance = 1 - similarity + np.fill_diagonal(distance, 0) + return cast(npt.NDArray[np.floating], distance) + + elif metric == "cosine": + if np.iscomplexobj(X): + raise ValueError("Cosine distance not compatible with complex data") + + # Use scipy's cosine distance + distances = pdist(X.T, metric='cosine') + return cast(npt.NDArray[np.floating], squareform(distances)) + + else: + raise ValueError(f"Distance metric '{metric}' not recognized") + + +def similarity_matrix( + X: npt.NDArray[np.floating], + distance_method: Literal["euclidean", "cosine", "mac"] = "euclidean", + graph_method: Literal["full", "knn", "mknn", "epsilon"] = "full", + **kwargs +) -> npt.NDArray[np.floating]: + """Compute similarity matrix for spectral clustering. + + Args: + X: Data matrix where each column is a feature vector + distance_method: Distance/similarity method + graph_method: Graph construction method + **kwargs: Additional parameters: + - sigma: For euclidean similarity (required) + - k: For knn/mknn methods (required) + - epsilon: For epsilon method (required) + + Returns: + Symmetric similarity matrix + """ + d, n = X.shape + + print(f"Processing {n} features (dimension {d}, {'complex' if np.iscomplexobj(X) else 'real'})") + + # Validate parameters + if distance_method == "euclidean" and "sigma" not in kwargs: + raise ValueError("sigma > 0 required for euclidean similarity") + if graph_method in ["knn", "mknn"] and "k" not in kwargs: + raise ValueError("k > 0 required for kNN methods") + if graph_method == "epsilon" and "epsilon" not in kwargs: + raise ValueError("epsilon > 0 required for epsilon method") + + if graph_method == "full": + if distance_method == "mac": + # Optimized vectorized MAC similarity + print("Using optimized full/MAC computation") + norms = np.sqrt(np.sum(X.conj() * X, axis=0, keepdims=True)) + norms[norms == 0] = 1 + X_norm = X / norms + S = np.abs(X_norm.conj().T @ X_norm) ** 2 + np.fill_diagonal(S, 0.0) + return cast(npt.NDArray[np.floating], S) + + elif distance_method == "euclidean": + sigma = kwargs["sigma"] + # Vectorized Euclidean similarity + distances = pdist(X.T, metric='euclidean') + dist_matrix = squareform(distances) + S = np.exp(-dist_matrix**2 / (2 * sigma**2)) + np.fill_diagonal(S, 0.0) + return cast(npt.NDArray[np.floating], S) + + elif distance_method == "cosine": + if np.iscomplexobj(X): + raise ValueError("Cosine similarity not compatible with complex data") + # Use cosine similarity (1 - cosine distance) + distances = pdist(X.T, metric='cosine') + dist_matrix = squareform(distances) + S = 1 - dist_matrix + np.fill_diagonal(S, 0.0) + return cast(npt.NDArray[np.floating], S) + + elif graph_method in ["knn", "mknn"]: + k = kwargs["k"] + # Remove k from kwargs to avoid conflicts + knn_kwargs = {key: value for key, value in kwargs.items() if key != "k"} + return _knn_similarity_matrix(X, distance_method, k, mutual=(graph_method == "mknn"), **knn_kwargs) + + elif graph_method == "epsilon": + epsilon = kwargs["epsilon"] + # First compute full similarity matrix + full_S = similarity_matrix(X, distance_method, "full", **kwargs) + # Apply epsilon threshold + threshold = 1 - epsilon + S = np.where(full_S >= threshold, full_S, 0) + return cast(npt.NDArray[np.floating], S) + + else: + raise ValueError(f"Graph method '{graph_method}' not recognized") + + +def _knn_similarity_matrix( + X: npt.NDArray[np.floating], + distance_method: Literal["euclidean", "cosine", "mac"], + k: int, + mutual: bool = False, + **kwargs +) -> npt.NDArray[np.floating]: + """Construct k-nearest neighbors similarity matrix.""" + n = X.shape[1] + + # First compute full similarity matrix efficiently + full_S = similarity_matrix(X, distance_method, "full", **kwargs) + + # Find k-nearest neighbors for each point + # Sort similarities in descending order to get top k + sorted_indices = np.argsort(-full_S, axis=1) + knn_indices = sorted_indices[:, :k] + + # Construct sparse similarity matrix + S = np.zeros((n, n)) + + for i in range(n): + neighbors_i = knn_indices[i] + for j in neighbors_i: + if i != j: + if mutual: + # Mutual kNN: check if i is also in j's neighborhood + neighbors_j = knn_indices[j] + if i in neighbors_j: + S[i, j] = full_S[i, j] + S[j, i] = full_S[i, j] + else: + # Regular kNN: set both directions to ensure symmetry + S[i, j] = full_S[i, j] + S[j, i] = full_S[i, j] + + return cast(npt.NDArray[np.floating], S) + + +def kmeans_clustering( + X: npt.NDArray[np.floating], + k: int, + n_trials: int = 10, + max_iter: int = 300, + random_state: Optional[int] = None +) -> ClusteringResult: + """K-means clustering with multiple random initializations. + + Args: + X: Data matrix where each column is a feature vector + k: Number of clusters + n_trials: Number of random initialization trials + max_iter: Maximum iterations per trial + random_state: Random seed for reproducibility + + Returns: + ClusteringResult with best clustering from all trials + """ + if k < 2: + raise ValueError("Number of clusters should be at least 2") + if n_trials < 1: + raise ValueError("At least 1 trial is required") + + if random_state is not None: + np.random.seed(random_state) + + d, n = X.shape + best_cost = float('inf') + best_result: Optional[ClusteringResult] = None + + for trial in range(n_trials): + # Random initialization from actual data points + initial_indices = np.random.choice(n, size=k, replace=False) + centroids = X[:, initial_indices].copy() + + result = _kmeans_single_trial(X, centroids, max_iter) + + if result.cost is not None and result.cost < best_cost: + best_cost = result.cost + best_result = result + + if best_result is None: + raise RuntimeError("K-means failed to converge in all trials") + + return best_result + + +def _kmeans_single_trial( + X: npt.NDArray[np.floating], + initial_centroids: npt.NDArray[np.floating], + max_iter: int +) -> ClusteringResult: + """Single trial of k-means clustering.""" + d, n = X.shape + k = initial_centroids.shape[1] + + centroids = initial_centroids.copy() + labels = np.zeros(n, dtype=np.int32) + + old_cost = float('inf') + total_cost = 0.0 + iteration = 0 + + for iteration in range(max_iter): + # Vectorized distance computation and assignment + distances = np.zeros((n, k)) + for j in range(k): + diff = X - centroids[:, j:j+1] + distances[:, j] = np.sum(np.real(diff.conj() * diff), axis=0) + + labels = np.argmin(distances, axis=1).astype(np.int32) + + # Update centroids and compute cost + total_cost = 0.0 + for j in range(k): + cluster_mask = labels == j + if np.any(cluster_mask): + cluster_points = X[:, cluster_mask] + centroids[:, j] = np.mean(cluster_points, axis=1) + + # Compute within-cluster sum of squares + diff = cluster_points - centroids[:, j:j+1] + total_cost += np.sum(np.real(diff.conj() * diff)) + + # Check convergence + if total_cost >= old_cost: + break + old_cost = total_cost + + return ClusteringResult( + labels=labels, + centroids=centroids, + cost=total_cost, + iterations=iteration + 1 + ) + + +def spectral_clustering( + similarity_matrix: npt.NDArray[np.floating], + k: Union[int, list], + method: Literal["standard", "shi", "ng"] = "standard", + n_trials: int = 10, + auto_select: bool = False +) -> ClusteringResult: + """Spectral clustering using various Laplacian methods. + + Args: + similarity_matrix: Symmetric similarity/adjacency matrix + k: Number of clusters (or list of k values, or 0 for auto-selection) + method: Laplacian method ('standard', 'shi', 'ng') + n_trials: Number of k-means trials in final stage + auto_select: Whether to auto-select k using eigengap heuristic + + Returns: + ClusteringResult with spectral clustering results + """ + if isinstance(k, (list, np.ndarray)): + if len(k) > 1: + raise NotImplementedError("Multiple k values not yet implemented") + k = k[0] if len(k) == 1 else 0 + + if k == 0: + auto_select = True + + n = similarity_matrix.shape[0] + if similarity_matrix.shape[1] != n: + raise ValueError("Similarity matrix must be square") + + # Ensure zero diagonal + W = similarity_matrix.copy() + if np.linalg.norm(np.diag(W)) != 0: + warnings.warn("Similarity matrix has nonzero diagonal (ignored)") + np.fill_diagonal(W, 0) + + # Compute degree matrix + degrees = np.sum(W, axis=1) + if np.any(degrees <= 0): + raise ValueError("All degree matrix entries must be positive") + + # Solve eigenvalue problem based on method + if method == "standard": + # Standard Laplacian: L = D - W + D = np.diag(degrees) + L = D - W + eigenvalues, eigenvectors = eigh(L) + + elif method in ["shi", "ng"]: + # Both use normalized Laplacian: I - D^(-1/2) W D^(-1/2) + D_inv_sqrt = np.diag(1.0 / np.sqrt(degrees)) + L_norm = np.eye(n) - D_inv_sqrt @ W @ D_inv_sqrt + eigenvalues, eigenvectors = eigh(L_norm) + + else: + raise ValueError(f"Laplacian method '{method}' not recognized") + + # Auto-select k using eigengap heuristic + if auto_select: + k = _eigengap_heuristic(eigenvalues) + if k < 2: + raise ValueError("Auto-selected k < 2, clustering not possible") + + if k < 2: + raise ValueError("Number of clusters should be at least 2") + + # Prepare features for k-means + if method == "standard": + features = eigenvectors[:, :k] + elif method == "shi": + D_inv_sqrt = np.diag(1.0 / np.sqrt(degrees)) + features = D_inv_sqrt @ eigenvectors[:, :k] + elif method == "ng": + # Row-normalize the eigenvectors + features = eigenvectors[:, :k] + row_norms = np.linalg.norm(features, axis=1, keepdims=True) + row_norms[row_norms == 0] = 1 # Avoid division by zero + features = features / row_norms + + # Final k-means clustering + kmeans_result = kmeans_clustering(features.T, k, n_trials) + + return ClusteringResult( + labels=kmeans_result.labels, + centroids=kmeans_result.centroids, + cost=kmeans_result.cost, + eigenvalues=eigenvalues, + iterations=kmeans_result.iterations + ) + + +def _eigengap_heuristic(eigenvalues: npt.NDArray[np.floating]) -> int: + """Auto-select number of clusters using eigengap heuristic.""" + n_eval = min(len(eigenvalues) // 2, len(eigenvalues)) + if n_eval < 2: + return 2 + + gaps = np.diff(eigenvalues[:n_eval]) + k = int(np.argmax(gaps)) + 1 + return max(k, 2) + + +def medoid_clustering( + X: npt.NDArray[np.floating], + labels: npt.NDArray[np.int32], + metric: Literal["euclidean", "cosine", "mac"] = "euclidean" +) -> npt.NDArray[np.int32]: + """Find cluster medoids given clustering labels. + + Args: + X: Data matrix where each column is a feature vector + labels: Cluster labels for each data point + metric: Distance metric for medoid computation + + Returns: + Array of medoid indices for each cluster + """ + unique_labels = np.unique(labels[labels > 0]) + num_clusters = len(unique_labels) + medoid_indices = np.zeros(num_clusters, dtype=np.int32) + + for i, cluster_id in enumerate(unique_labels): + cluster_points = np.where(labels == cluster_id)[0] + if len(cluster_points) == 0: + raise ValueError(f"Empty cluster #{cluster_id}") + + medoid_idx = _find_medoid(X[:, cluster_points], metric) + medoid_indices[i] = cluster_points[medoid_idx] + + return medoid_indices + + +def _find_medoid(X: npt.NDArray[np.floating], metric: Literal["euclidean", "cosine", "mac"]) -> int: + """Find medoid (point with smallest average distance to all others).""" + n = X.shape[1] + D = distance_matrix(X, metric) + + # Find point with minimum sum of distances + costs = np.sum(D, axis=1) + return int(np.argmin(costs)) + + +def mac_value(v: npt.NDArray, w: npt.NDArray) -> float: + """Compute Modal Assurance Criterion (MAC) value between two vectors.""" + return DistanceMetric.mac_similarity(v, w) + + +def trim_cluster_mac( + X: npt.NDArray[np.floating], + labels: npt.NDArray[np.int32], + medoid_indices: npt.NDArray[np.int32], + threshold: float +) -> npt.NDArray[np.int32]: + """Remove vectors from clusters if MAC with medoid is below threshold. + + Args: + X: Data matrix where each column is a feature vector + labels: Current cluster labels + medoid_indices: Indices of cluster medoids + threshold: MAC threshold (0 < threshold < 1) + + Returns: + Updated cluster labels (demoted points get label 0) + """ + if threshold <= 0 or threshold >= 1: + raise ValueError("Threshold must satisfy 0 < threshold < 1") + + updated_labels = labels.copy() + total_removed = 0 + num_clusters = len(medoid_indices) + + for cluster_id in range(1, num_clusters + 1): + cluster_points = np.where(labels == cluster_id)[0] + if len(cluster_points) == 0: + continue + + medoid_idx = medoid_indices[cluster_id - 1] + medoid_shape = X[:, medoid_idx] + + # Vectorized MAC computation for all points in cluster + cluster_shapes = X[:, cluster_points] + # Compute MAC for each column + numerators = np.abs(medoid_shape.conj().T @ cluster_shapes) ** 2 + medoid_norm_sq = np.real(medoid_shape.conj().T @ medoid_shape) + cluster_norms_sq = np.sum(cluster_shapes.conj() * cluster_shapes, axis=0) + denominators = medoid_norm_sq * cluster_norms_sq + + # Avoid division by zero + mac_values = np.where(denominators > 0, numerators / denominators, 0.0) + mac_values = np.real(mac_values) + + # Find points below threshold + below_threshold = mac_values < threshold + removed_points = cluster_points[below_threshold] + updated_labels[removed_points] = 0 + + removed_count = len(removed_points) + print(f"#{removed_count} vectors removed from cluster #{cluster_id}") + total_removed += removed_count + + print(f"Total of #{total_removed} vectors removed from clusters") + return updated_labels + + +# ============================================================================= +# Missing MATLAB Clustering Functions +# ============================================================================= + +def clus_similarity_matrix( + X: npt.NDArray, + distance_method: str = 'euclidean', + graph_method: str = 'full', + method_args: Optional[Dict] = None +) -> npt.NDArray[np.floating]: + """ + Calculate symmetric similarity matrix for clustering. + + This function calculates similarity matrices from feature data using various + distance and graph construction methods, following the MATLAB clus_similarity_matrix.m. + + Args: + X: Feature matrix where each column is a feature vector, shape (d, n) + distance_method: Distance method ('euclidean', 'cosine', 'mac') + graph_method: Graph construction method ('full', 'knn', 'mknn', 'epsilon') + method_args: Additional arguments for graph methods (e.g., {'k': 10} for knn) + + Returns: + Symmetric similarity matrix, shape (n, n) + + Example: + >>> X = np.random.randn(5, 100) # 100 features of dimension 5 + >>> S = clus_similarity_matrix(X, 'euclidean', 'knn', {'k': 10}) + >>> print(f"Similarity matrix shape: {S.shape}") + """ + if method_args is None: + method_args = {} + + d, n = X.shape + data_is_complex = np.iscomplexobj(X) + + # Compute distance matrix + if distance_method == 'euclidean': + if data_is_complex: + # For complex data, use magnitude + X_real = np.abs(X) + distances = np.zeros((n, n)) + for i in range(n): + for j in range(i, n): + dist = np.linalg.norm(X_real[:, i] - X_real[:, j]) + distances[i, j] = distances[j, i] = dist + else: + # Standard Euclidean distance + distances = np.zeros((n, n)) + for i in range(n): + for j in range(i, n): + dist = np.linalg.norm(X[:, i] - X[:, j]) + distances[i, j] = distances[j, i] = dist + + elif distance_method == 'cosine': + if data_is_complex: + raise NotImplementedError("Cosine similarity not implemented for complex data") + # Cosine distance = 1 - cosine similarity + from sklearn.metrics.pairwise import cosine_similarity + similarities = cosine_similarity(X.T) + distances = 1 - similarities + + elif distance_method == 'mac': + # Modal Assurance Criterion (1 - MAC) + distances = np.zeros((n, n)) + for i in range(n): + for j in range(i, n): + if data_is_complex: + numerator = np.abs(X[:, i].conj().T @ X[:, j])**2 + denominator = np.real((X[:, i].conj().T @ X[:, i]) * + (X[:, j].conj().T @ X[:, j])) + else: + numerator = (X[:, i].T @ X[:, j])**2 + denominator = (X[:, i].T @ X[:, i]) * (X[:, j].T @ X[:, j]) + + mac_value = numerator / denominator if denominator > 0 else 0 + distances[i, j] = distances[j, i] = 1 - mac_value + else: + raise ValueError(f"Unknown distance method: {distance_method}") + + # Convert distances to similarities + if distance_method == 'euclidean': + # Use Gaussian kernel + sigma = method_args.get('sigma', np.std(distances[distances > 0])) + similarities = np.exp(-distances**2 / (2 * sigma**2)) + else: + similarities = 1 - distances + + # Apply graph construction method + if graph_method == 'full': + # Use full similarity matrix + S = similarities + + elif graph_method == 'knn': + # k-nearest neighbors + k = method_args.get('k', 10) + S = np.zeros_like(similarities) + + for i in range(n): + # Find k nearest neighbors (excluding self) + neighbor_indices = np.argsort(distances[i, :])[:k+1] + neighbor_indices = neighbor_indices[neighbor_indices != i][:k] + S[i, neighbor_indices] = similarities[i, neighbor_indices] + + elif graph_method == 'mknn': + # Mutual k-nearest neighbors + k = method_args.get('k', 10) + knn_graph = np.zeros_like(similarities) + + # First build knn graph + for i in range(n): + neighbor_indices = np.argsort(distances[i, :])[:k+1] + neighbor_indices = neighbor_indices[neighbor_indices != i][:k] + knn_graph[i, neighbor_indices] = 1 + + # Make it mutual (symmetric) + mutual_graph = knn_graph * knn_graph.T + S = similarities * mutual_graph + + elif graph_method == 'epsilon': + # Epsilon neighborhood + epsilon = method_args.get('epsilon', np.median(distances[distances > 0])) + S = similarities * (distances <= epsilon) + + else: + raise ValueError(f"Unknown graph method: {graph_method}") + + # Ensure symmetry + S = (S + S.T) / 2 + + return S + + +def clus_distance_matrix( + X: npt.NDArray, + distance_method: str = 'euclidean' +) -> npt.NDArray[np.floating]: + """ + Compute pairwise distance matrix between feature vectors. + + Args: + X: Feature matrix where each column is a feature vector, shape (d, n) + distance_method: Distance method ('euclidean', 'manhattan', 'mac') + + Returns: + Distance matrix, shape (n, n) + """ + d, n = X.shape + distances = np.zeros((n, n)) + + if distance_method == 'euclidean': + for i in range(n): + for j in range(i, n): + dist = np.linalg.norm(X[:, i] - X[:, j]) + distances[i, j] = distances[j, i] = dist + + elif distance_method == 'manhattan': + for i in range(n): + for j in range(i, n): + dist = np.sum(np.abs(X[:, i] - X[:, j])) + distances[i, j] = distances[j, i] = dist + + elif distance_method == 'mac': + # MAC-based distance (1 - MAC) + for i in range(n): + for j in range(i, n): + if np.iscomplexobj(X): + numerator = np.abs(X[:, i].conj().T @ X[:, j])**2 + denominator = np.real((X[:, i].conj().T @ X[:, i]) * + (X[:, j].conj().T @ X[:, j])) + else: + numerator = (X[:, i].T @ X[:, j])**2 + denominator = (X[:, i].T @ X[:, i]) * (X[:, j].T @ X[:, j]) + + mac_value = numerator / denominator if denominator > 0 else 0 + distances[i, j] = distances[j, i] = 1 - mac_value + else: + raise ValueError(f"Unknown distance method: {distance_method}") + + return distances + + +def spclus_knn_similarity_matrix( + X: npt.NDArray, + k: int = 10, + sigma: Optional[float] = None +) -> npt.NDArray[np.floating]: + """ + Create k-NN similarity matrix for spectral clustering. + + Args: + X: Feature matrix, shape (d, n) + k: Number of nearest neighbors + sigma: Gaussian kernel width (auto-estimated if None) + + Returns: + k-NN similarity matrix, shape (n, n) + """ + return clus_similarity_matrix(X, 'euclidean', 'knn', {'k': k, 'sigma': sigma}) + + +def spclus_spectral( + similarity_matrix: npt.NDArray[np.floating], + n_clusters: int, + method: str = 'normalized' +) -> npt.NDArray[np.int32]: + """ + Spectral clustering using eigendecomposition of similarity matrix. + + Args: + similarity_matrix: Symmetric similarity matrix, shape (n, n) + n_clusters: Number of clusters + method: Spectral clustering variant ('normalized', 'unnormalized') + + Returns: + Cluster labels, shape (n,) + """ + from sklearn.cluster import SpectralClustering + + # Use sklearn's implementation + spectral = SpectralClustering( + n_clusters=n_clusters, + affinity='precomputed', + assign_labels='kmeans', + n_init=10 + ) + + labels = spectral.fit_predict(similarity_matrix) + + # Convert to 1-based indexing to match MATLAB + return labels + 1 + + +def clus_krnn_enhance( + X: npt.NDArray, + similarity_matrix: npt.NDArray[np.floating], + k: int = 5, + enhancement_factor: float = 2.0 +) -> npt.NDArray[np.floating]: + """ + Enhance similarity matrix using k-reciprocal nearest neighbors. + + This function enhances the similarity matrix by identifying reciprocal + nearest neighbors and boosting their similarity values. + + Args: + X: Feature matrix, shape (d, n) + similarity_matrix: Input similarity matrix, shape (n, n) + k: Number of nearest neighbors to consider + enhancement_factor: Factor by which to enhance reciprocal similarities + + Returns: + Enhanced similarity matrix, shape (n, n) + """ + n = similarity_matrix.shape[0] + enhanced_matrix = similarity_matrix.copy() + + # Compute distance matrix for neighbor finding + distances = 1 - similarity_matrix + + # Find reciprocal nearest neighbors + for i in range(n): + # Find k nearest neighbors of point i + neighbors_i = np.argsort(distances[i, :])[:k+1] + neighbors_i = neighbors_i[neighbors_i != i][:k] + + for j in neighbors_i: + # Check if i is also among k nearest neighbors of j + neighbors_j = np.argsort(distances[j, :])[:k+1] + neighbors_j = neighbors_j[neighbors_j != j][:k] + + if i in neighbors_j: + # i and j are reciprocal neighbors - enhance similarity + enhanced_matrix[i, j] *= enhancement_factor + enhanced_matrix[j, i] *= enhancement_factor + + return enhanced_matrix \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/data_extraction.py b/src/tokeye/eigspec/utils/data_extraction.py new file mode 100644 index 0000000..5ab29ef --- /dev/null +++ b/src/tokeye/eigspec/utils/data_extraction.py @@ -0,0 +1,376 @@ +""" +Data extraction and collection utilities for eigspec package. + +This module provides functions for extracting and organizing analysis results: +- extract_ptrefs: Extract reference points from analysis results +- collect_rep_data: Harvest features from analysis results +- Data organization and formatting utilities + +Based on the MATLAB eigspec toolbox data collection functions. +""" + +from typing import Optional, List, Tuple, Dict, Any, Union +import numpy as np +import numpy.typing as npt +from dataclasses import dataclass + + +@dataclass +class PointReference: + """Single point reference extracted from analysis results.""" + query_tifr: Tuple[float, float] # [time_ms, freq_kHz] + shapevector: npt.NDArray[np.complex128] + rms: float + frequency: float # Hz + radius: float + centre_time: float # seconds + + +@dataclass +class PointReferences: + """Collection of point references with metadata.""" + modes: List[PointReference] + sensor_coordinates: npt.NDArray[np.floating] + Ts: float # sampling period in seconds + block_params: Tuple[int, int] + analysis_params: Tuple[int, ...] + thresholds: Tuple[float, float] + + +def extract_ptrefs( + analysis_result: Any, + query_points: npt.NDArray[np.floating], + sensor_coordinates: npt.NDArray[np.floating], + warning_distance: float = 5.0 +) -> PointReferences: + """ + Extract reference points from analysis results. + + Python port of MATLAB extract_ptrefs.m that finds the closest modal + features to provided query points in (time, frequency) space. + + Args: + analysis_result: Analysis result structure with .L blocks + query_points: Query points [time_ms, freq_kHz], shape (n_refs, 2) + sensor_coordinates: Sensor coordinate matrix + warning_distance: Distance threshold for warnings + + Returns: + PointReferences object with extracted reference points + """ + if not hasattr(analysis_result, 'L') or not analysis_result.L: + raise ValueError("Analysis result must have non-empty L field") + + n_refs = query_points.shape[0] + if query_points.shape[1] != 2: + raise ValueError("Query points must have shape (n_refs, 2) for [time, freq]") + + # Collect all available features + XTRF, XRMS, XS = collect_rep_data(analysis_result) + + if XTRF.shape[0] == 0: + raise ValueError("No modal features found in analysis results") + + # Convert frequency to kHz for comparison + XTRF_search = XTRF.copy() + XTRF_search[:, 2] = XTRF_search[:, 2] / 1e3 # Convert Hz to kHz + + # Extract metadata + first_block = analysis_result.L[0] + Ts = getattr(first_block, 'Ts', -1) + block_params = getattr(analysis_result, 'bss', (0, 0)) + analysis_params = getattr(analysis_result, 'rfpn', ()) + thresholds = getattr(analysis_result, 'thresh', (0.0, 0.0)) + + # Find closest matches for each query point + extracted_modes = [] + + for jj in range(n_refs): + query_time, query_freq = query_points[jj] + + # Find closest point in (time, frequency) space + distances_sq = ( + (XTRF_search[:, 0] - query_time)**2 + + (XTRF_search[:, 2] - query_freq)**2 + ) + + closest_idx = np.argmin(distances_sq) + distance = np.sqrt(distances_sq[closest_idx]) + + closest_time = XTRF_search[closest_idx, 0] + closest_freq = XTRF_search[closest_idx, 2] + + print(f"Query {jj+1}: closest to [time,freq]=[{query_time:.6f},{query_freq:.6f}] " + f"found at index {closest_idx+1}/{len(XTRF)} " + f"[time,freq]=[{closest_time:.6f},{closest_freq:.6f}] (ms,kHz)") + + if distance > warning_distance: + print(f"(Warning: distance to query point is large, d={distance:.3f})") + + # Create point reference + mode_ref = PointReference( + query_tifr=(query_time, query_freq), + shapevector=XS[:, closest_idx], + rms=XRMS[closest_idx], + frequency=XTRF[closest_idx, 2], # Hz + radius=XTRF[closest_idx, 1], + centre_time=XTRF[closest_idx, 0] / 1e3 # Convert ms to seconds + ) + + extracted_modes.append(mode_ref) + + return PointReferences( + modes=extracted_modes, + sensor_coordinates=sensor_coordinates, + Ts=Ts, + block_params=block_params, + analysis_params=analysis_params, + thresholds=thresholds + ) + + +def collect_rep_data(analysis_result: Any) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.complex128]]: + """ + Harvest features from analysis results and store in arrays. + + Python port of MATLAB collect_rep_data.m that extracts all modal + features from block-based analysis results. + + Args: + analysis_result: Analysis result structure with .L blocks + + Returns: + XTRF: Feature matrix [time_ms, damping, frequency_Hz], shape (n_modes, 3) + XRMS: RMS values for each mode, shape (n_modes,) + XS: Complex mode shape vectors, shape (n_channels, n_modes) + """ + if not hasattr(analysis_result, 'L') or not analysis_result.L: + return np.array([]).reshape(0, 3), np.array([]), np.array([]).reshape(0, 0) + + # Determine number of channels + first_block = analysis_result.L[0] + if hasattr(analysis_result, 'iisubset'): + M = len(analysis_result.iisubset) + elif hasattr(first_block, 'Psi'): + M = first_block.Psi.shape[1] if hasattr(first_block.Psi, 'shape') else 1 + else: + M = 1 # Default fallback + + Ts = getattr(first_block, 'Ts', -1) + n_blocks = len(analysis_result.L) + + # Count total modes and collect block times + T = np.zeros(n_blocks) + total_modes = 0 + + for jj, block in enumerate(analysis_result.L): + if hasattr(block, 'mrep') and hasattr(block.mrep, 'imode'): + total_modes += len(block.mrep.imode) + + if hasattr(block, 'centre_t'): + T[jj] = block.centre_t + else: + T[jj] = jj # Default time indexing + + # Initialize output arrays + if total_modes == 0: + return np.array([]).reshape(0, 3), np.array([]), np.array([]).reshape(0, 0) + + XTRF = np.zeros((total_modes, 3)) # [time_ms, damping, frequency_Hz] + XRMS = np.zeros(total_modes) + XS = np.zeros((M, total_modes), dtype=np.complex128) + + # Extract features from each block + mode_idx = 0 + + for jj, block in enumerate(analysis_result.L): + if not (hasattr(block, 'mrep') and hasattr(block.mrep, 'imode')): + continue + + n_modes_block = len(block.mrep.imode) + + for mm in range(n_modes_block): + # Get modal shape vector + C = _get_c_shape(block, mm, M) + + # Get frequency and damping + radius, frequency = _get_freq(block, mm, Ts) + + # Store data + XTRF[mode_idx, :] = [1e3 * T[jj], radius, frequency] # time in ms, freq in Hz + XS[:, mode_idx] = C + + # Get RMS if available + if (hasattr(block, 'drep') and hasattr(block.drep, 'drms') and + mm < len(block.drep.drms)): + XRMS[mode_idx] = block.drep.drms[mm] + else: + XRMS[mode_idx] = np.linalg.norm(C) # Default to vector norm + + mode_idx += 1 + + return XTRF, XRMS, XS + + +def _get_freq(block: Any, mode_idx: int, Ts: float) -> Tuple[float, float]: + """Extract frequency and damping from block data.""" + if not (hasattr(block, 'mrep') and hasattr(block.mrep, 'm0') and + hasattr(block.mrep.m0, 'lambda')): + return 0.0, 0.0 + + if mode_idx >= len(block.mrep.imode): + return 0.0, 0.0 + + mode_number = block.mrep.imode[mode_idx] + # 'lambda' is a keyword: attribute access must go through getattr + eigenvalues = getattr(block.mrep.m0, 'lambda') + if mode_number >= len(eigenvalues): + return 0.0, 0.0 + + eigenvalue = eigenvalues[mode_number] + + # Extract radius (damping) and frequency + radius = np.abs(eigenvalue) + frequency = np.angle(eigenvalue) + + # Convert to Hz if sampling period is provided + if Ts > 0: + frequency = frequency / (2 * np.pi * Ts) + + return float(radius), float(frequency) + + +def _get_c_shape(block: Any, mode_idx: int, n_channels: int) -> npt.NDArray[np.complex128]: + """Extract complex shape vector from block data.""" + if not (hasattr(block, 'mrep') and hasattr(block.mrep, 'm0')): + return np.zeros(n_channels, dtype=np.complex128) + + if mode_idx >= len(block.mrep.imode): + return np.zeros(n_channels, dtype=np.complex128) + + # Try to get shape from various possible locations + if hasattr(block.mrep.m0, 'shape'): + shapes = block.mrep.m0.shape + if shapes.size > 0: + if shapes.ndim == 2 and mode_idx < shapes.shape[1]: + return shapes[:, mode_idx] + elif shapes.ndim == 1 and mode_idx == 0: + return shapes + + # Try alternative shape storage + if hasattr(block, 'drep') and hasattr(block.drep, 'dhat'): + dhat = block.drep.dhat + if dhat.size > 0: + # Reconstruct complex shape from real/imaginary parts + if dhat.ndim == 2 and mode_idx * 2 + 1 < dhat.shape[1]: + real_part = dhat[:, mode_idx * 2] + imag_part = dhat[:, mode_idx * 2 + 1] + return real_part + 1j * imag_part + + # Default fallback + return np.zeros(n_channels, dtype=np.complex128) + + +def collect_prototype_traces( + analysis_result: Any, + point_references: PointReferences, + time_vector: npt.NDArray[np.floating], + signal_data: npt.NDArray[np.floating] +) -> Dict[str, npt.NDArray[np.floating]]: + """ + Collect prototype time traces for reference points. + + Args: + analysis_result: Analysis result structure + point_references: Reference points extracted from analysis + time_vector: Time vector for signal data + signal_data: Original signal data matrix + + Returns: + Dictionary with prototype traces for each reference point + """ + n_refs = len(point_references.modes) + prototype_traces = {} + + for i, mode_ref in enumerate(point_references.modes): + # Find time window around reference point + center_time = mode_ref.centre_time + frequency = mode_ref.frequency + + # Use several periods for prototype extraction + if frequency > 0: + period = 1.0 / frequency + window_duration = min(5 * period, 0.1) # 5 periods or 100ms max + else: + window_duration = 0.05 # 50ms default + + # Find indices for time window + time_mask = (np.abs(time_vector - center_time) <= window_duration / 2) + + if np.any(time_mask): + window_data = signal_data[time_mask, :] + window_time = time_vector[time_mask] - center_time # Relative time + + # Compute weighted projection using shape vector + shape_vec = mode_ref.shapevector + if len(shape_vec) == window_data.shape[1]: + # Project signal onto mode shape + prototype = np.real(window_data @ np.conj(shape_vec)) + + prototype_traces[f'mode_{i+1}'] = { + 'time': window_time, + 'signal': prototype, + 'reference_info': { + 'frequency': frequency, + 'rms': mode_ref.rms, + 'query_point': mode_ref.query_tifr + } + } + + return prototype_traces + + +def format_analysis_summary(analysis_result: Any) -> Dict[str, Any]: + """ + Create a summary of analysis results. + + Args: + analysis_result: Analysis result structure + + Returns: + Dictionary with analysis summary statistics + """ + summary = { + 'n_blocks': 0, + 'total_modes': 0, + 'frequency_range': [0.0, 0.0], + 'time_range': [0.0, 0.0], + 'processing_method': 'unknown' + } + + if not hasattr(analysis_result, 'L') or not analysis_result.L: + return summary + + # Basic counts + summary['n_blocks'] = len(analysis_result.L) + summary['processing_method'] = getattr(analysis_result, 'routine', 'unknown') + + # Collect features for statistics + XTRF, XRMS, XS = collect_rep_data(analysis_result) + + if XTRF.shape[0] > 0: + summary['total_modes'] = XTRF.shape[0] + summary['time_range'] = [float(XTRF[:, 0].min()), float(XTRF[:, 0].max())] + summary['frequency_range'] = [float(XTRF[:, 2].min()), float(XTRF[:, 2].max())] + summary['rms_range'] = [float(XRMS.min()), float(XRMS.max())] + summary['damping_range'] = [float(XTRF[:, 1].min()), float(XTRF[:, 1].max())] + + # Add parameter information if available + if hasattr(analysis_result, 'bss'): + summary['block_parameters'] = analysis_result.bss + if hasattr(analysis_result, 'rfpn'): + summary['analysis_parameters'] = analysis_result.rfpn + if hasattr(analysis_result, 'thresh'): + summary['thresholds'] = analysis_result.thresh + + return summary \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/fdm_analysis.py b/src/tokeye/eigspec/utils/fdm_analysis.py new file mode 100644 index 0000000..3c18f10 --- /dev/null +++ b/src/tokeye/eigspec/utils/fdm_analysis.py @@ -0,0 +1,435 @@ +""" +Finite Difference Method (FDM) analysis for eigspec package. + +This module provides FDM-based frequency detection algorithms that are fundamental +to the eigspec toolbox: +- fdm1d: Core 1D FDM frequency detection algorithm +- Related FDM utilities and helper functions + +Based on the MATLAB eigspec toolbox FDM implementations. +""" + +from typing import Optional, Tuple, Dict, Any +import numpy as np +import numpy.typing as npt +from scipy.linalg import eig, svd +import warnings + + +def fdm1d( + Y: npt.NDArray[np.floating], + d: int, + alpha: float = -1.0 +) -> Dict[str, Any]: + """ + FDM-type frequency detection for multichannel time-series data. + + Python port of MATLAB fdm1d.m that implements finite difference method + for finding frequency lists from multichannel time-series. Uses generalized + eigenvalue problems with regularization to detect dominant frequencies. + + Args: + Y: Signal data matrix where columns are channels, shape (n, m) + d: Desired eigenvalue problem dimension + alpha: Regularization parameter (negative for auto-selection) + + Returns: + Dictionary containing: + - eigk: Sorted eigenvalues (decreasing plausibility) + - errk: Error estimates for each eigenvalue + - tallness: Data matrix tallness ratio + - p: Snapshot length + - alpha: Used regularization parameter + """ + default_epscut = 1e-12 + minimum_snapshot = 10 + + n, m = Y.shape + p = n - d + + if p < minimum_snapshot: + raise ValueError("Snapshot length is too small") + + # Check tallness requirement + tallness = m * p / (n - p + 1) + if tallness < 2: + raise ValueError("Code requires data matrix to be significantly tall") + + # Build data matrix D + nd = n - p + 1 + D = np.zeros((m * p, nd)) + + for tt in range(p, n): # MATLAB: (p+1):(n+1), but 0-indexed + jj = tt - p + # Extract p consecutive samples in reverse order for each channel + block = Y[tt:tt-p:-1, :].T.ravel() # Equivalent to MATLAB reshape + D[:, jj] = block + + # Construct shifted matrices for generalized eigenvalue problem + S = D[:, :-1] # D(:,1:(nd-1)) + R = D[:, 1:] # D(:,2:nd) + + nm = S.shape[1] + + # Auto-select regularization parameter if needed + if alpha < 0: + # Based on eigenvalues of S'*S following Werner&Cary approach + StS = S.T @ S + eig_vals = np.linalg.eigvals(StS) + eig_vals = np.sort(np.real(eig_vals))[::-1] # Descending order + + # Use median of eigenvalues as regularization (heuristic) + alpha = float(np.median(eig_vals)) * 1e-6 + + if alpha < default_epscut: + alpha = default_epscut + + # Solve regularized generalized eigenvalue problem + # (R'*R + alpha*I) * v = lambda * (S'*S + alpha*I) * v + A = R.T @ R + alpha * np.eye(nm) + B = S.T @ S + alpha * np.eye(nm) + + try: + eigenvals, eigenvecs = eig(A, B) + except np.linalg.LinAlgError as e: + warnings.warn(f"Eigenvalue computation failed: {e}") + return { + 'eigk': np.array([]), + 'errk': np.array([]), + 'tallness': tallness, + 'p': p, + 'alpha': alpha + } + + # Convert to complex eigenvalues (frequencies) + frequencies = np.log(eigenvals) / (1j * 2 * np.pi) + + # Compute error estimates (based on residual norms) + error_estimates = np.zeros(len(frequencies)) + + for i, (eigval, eigvec) in enumerate(zip(eigenvals, eigenvecs.T)): + if np.abs(eigval) > default_epscut: + # Compute residual for error estimate + residual_A = A @ eigvec - eigval * (B @ eigvec) + residual_norm = np.linalg.norm(residual_A) + + # Normalize by eigenvalue magnitude + error_estimates[i] = residual_norm / max(np.abs(eigval), default_epscut) + else: + error_estimates[i] = np.inf + + # Sort by error estimates (ascending = most plausible first) + sort_idx = np.argsort(error_estimates) + sorted_frequencies = frequencies[sort_idx] + sorted_errors = error_estimates[sort_idx] + + # Convert error estimates to confidence values (higher = better) + confidence = 1.0 / (1.0 + sorted_errors) + + return { + 'eigk': sorted_frequencies, + 'errk': confidence, + 'tallness': tallness, + 'p': p, + 'alpha': alpha, + 'raw_eigenvals': eigenvals[sort_idx], + 'eigenvecs': eigenvecs[:, sort_idx] + } + + +def rndspec( + time_vector: Optional[npt.NDArray[np.floating]], + signal_data: npt.NDArray[np.floating], + block_params: Tuple[int, int], + analysis_params: Tuple[int, ...], + thresholds: Tuple[float, float], + use_cca: bool = False +) -> Dict[str, Any]: + """ + Block-based compressed sampling SSI/AR/PCA analysis. + + Python port of MATLAB rndspec.m that performs block-based spectral analysis + using random projection and various identification methods. + + Args: + time_vector: Time vector or None for sample indices + signal_data: Signal data matrix, shape (N, M) + block_params: [block_size, block_stride] + analysis_params: [reduced_dim, future, past, order1, order2] for SSI + or [reduced_dim, past, order1, order2] for AR/PCA + thresholds: [MAC_threshold, DST_threshold] + use_cca: Whether to use CCA/CVA in SSI subprogram + + Returns: + Dictionary with analysis results for each block + """ + block_size, block_stride = block_params + N, M = signal_data.shape + + # Handle time vector + if time_vector is None or (len(time_vector) == 1 and time_vector[0] == 0): + time_vector = np.arange(1, N + 1) + Ts = -1 + else: + time_vector = time_vector.ravel() + if len(time_vector) != N: + raise ValueError("Length of time vector does not match signal data") + Ts = time_vector[1] - time_vector[0] + + # Block analysis setup + t1_vec = np.arange(0, N - block_size + 1, block_stride) + n_blocks = len(t1_vec) + + # Initialize results + block_results = [] + + for jj in range(n_blocks): + t1 = t1_vec[jj] + t2 = t1 + block_size + + # Extract block data + block_data = signal_data[t1:t2, :] + + # Process block using appropriate method + if len(analysis_params) == 5: + # SSI mode + block_result = _process_ssi_block( + block_data, analysis_params, thresholds, use_cca, Ts + ) + else: + # AR/PCA mode + block_result = _process_arpca_block( + block_data, analysis_params, thresholds, Ts + ) + + # Add timing and block info + block_result['t1t2'] = [t1, t2] + block_result['centre_t'] = (time_vector[t1] + time_vector[t2-1]) / 2 + block_result['filter_t'] = time_vector[t2-1] + block_result['Ts'] = Ts + + block_results.append(block_result) + + return { + 'L': block_results, + 'bss': block_params, + 'rfpn': analysis_params, + 'thresh': thresholds, + 'routine': 'rndspec' + } + + +def pcaspecx( + time_vector: Optional[npt.NDArray[np.floating]], + signal_data: npt.NDArray[np.floating], + block_params: Tuple[int, int], + analysis_params: Tuple[int, int, int, int, int], + thresholds: Tuple[float, float], + k_folds: int = 8 +) -> Dict[str, Any]: + """ + Block-based PCA-projected SSI analysis with k-fold pruning. + + Python port of MATLAB pcaspecx.m that uses PCA projection with k-fold + pruning to reduce outlier effects on the projection. + + Args: + time_vector: Time vector or None for sample indices + signal_data: Signal data matrix, shape (N, M) + block_params: [block_size, block_stride] + analysis_params: [reduced_dim, future, past, order1, order2] + thresholds: [MAC_threshold, DST_threshold] + k_folds: Number of folds for pruning (default 8) + + Returns: + Dictionary with analysis results for each block + """ + block_size, block_stride = block_params + N, M = signal_data.shape + + # Handle time vector + if time_vector is None or (len(time_vector) == 1 and time_vector[0] == 0): + time_vector = np.arange(1, N + 1) + Ts = -1 + else: + time_vector = time_vector.ravel() + if len(time_vector) != N: + raise ValueError("Length of time vector does not match signal data") + Ts = time_vector[1] - time_vector[0] + + # Block analysis setup + t1_vec = np.arange(0, N - block_size + 1, block_stride) + n_blocks = len(t1_vec) + + # Initialize results + block_results = [] + + for jj in range(n_blocks): + t1 = t1_vec[jj] + t2 = t1 + block_size + + # Extract block data + block_data = signal_data[t1:t2, :] + + # Process block with PCA projection and k-fold pruning + block_result = _process_pca_block( + block_data, analysis_params, thresholds, k_folds, 'mcd' + ) + + # Add timing and block info + block_result['t1t2'] = [t1, t2] + block_result['centre_t'] = (time_vector[t1] + time_vector[t2-1]) / 2 + block_result['filter_t'] = time_vector[t2-1] + block_result['Ts'] = Ts + + block_results.append(block_result) + + return { + 'L': block_results, + 'bss': block_params, + 'rfpn': analysis_params, + 'thresh': thresholds, + 'k': k_folds, + 'routine': 'pcaspecx' + } + + +# Helper functions for block processing + +def _process_ssi_block( + block_data: npt.NDArray[np.floating], + params: Tuple[int, int, int, int, int], + thresholds: Tuple[float, float], + use_cca: bool, + Ts: float +) -> Dict[str, Any]: + """Process a single block using SSI method.""" + # Placeholder implementation - would need full SSI algorithm + from ..subspace_identification import covariance_driven_ssi + + reduced_dim, future, past, order1, order2 = params + + # Apply random projection if needed + if reduced_dim > 0 and reduced_dim < block_data.shape[1]: + projection = np.random.randn(reduced_dim, block_data.shape[1]) + if reduced_dim < 0: # Orthonormalize + projection = np.linalg.qr(projection.T)[0].T + block_data = (projection @ block_data.T).T + + # Run SSI + try: + ssi_result = covariance_driven_ssi( + block_data, + model_order=max(order1, order2), + block_rows=future + past + ) + + # Extract modes and format results + result = { + 'mrep': { + 'imode': np.arange(len(ssi_result.natural_frequencies)), + 'm0': { + 'lambda': ssi_result.discrete_eigenvalues, + 'shape': ssi_result.mode_shapes + } + }, + 'method': 'ssi', + 'Psi': projection if reduced_dim > 0 else np.eye(block_data.shape[1]) + } + + except Exception as e: + # Return empty result if SSI fails + result = { + 'mrep': {'imode': np.array([]), 'm0': {'lambda': np.array([]), 'shape': np.array([])}}, + 'method': 'ssi_failed', + 'error': str(e) + } + + return result + + +def _process_arpca_block( + block_data: npt.NDArray[np.floating], + params: Tuple[int, int, int, int], + thresholds: Tuple[float, float], + Ts: float +) -> Dict[str, Any]: + """Process a single block using AR/PCA method.""" + from ..autoregressive_pca import arpca + + reduced_dim, past, order1, order2 = params + + # Apply random projection if needed + projection = np.eye(block_data.shape[1]) + if reduced_dim > 0 and reduced_dim < block_data.shape[1]: + projection = np.random.randn(reduced_dim, block_data.shape[1]) + if reduced_dim < 0: # Orthonormalize + projection = np.linalg.qr(projection.T)[0].T + block_data = (projection @ block_data.T).T + + # Run AR/PCA + try: + arpca_result = arpca(block_data, max(order1, order2), reduced_dim or block_data.shape[1]) + + result = { + 'mrep': { + 'imode': np.arange(len(arpca_result.model.eigenvalues)), + 'm0': { + 'lambda': arpca_result.model.eigenvalues, + 'shape': arpca_result.model.mode_shapes + } + }, + 'method': 'arpca', + 'Psi': projection + } + + except Exception as e: + result = { + 'mrep': {'imode': np.array([]), 'm0': {'lambda': np.array([]), 'shape': np.array([])}}, + 'method': 'arpca_failed', + 'error': str(e) + } + + return result + + +def _process_pca_block( + block_data: npt.NDArray[np.floating], + params: Tuple[int, int, int, int, int], + thresholds: Tuple[float, float], + k_folds: int, + criterion: str +) -> Dict[str, Any]: + """Process a single block using PCA projection with k-fold pruning.""" + from ..matlab_utilities import kfoldcov + + reduced_dim, future, past, order1, order2 = params + + # Apply k-fold pruning for robust PCA + try: + idx, mu, R, _ = kfoldcov(block_data.T, k_folds, criterion) + pruned_data = block_data[idx, :] + + # Compute PCA projection + if reduced_dim > 0 and reduced_dim < pruned_data.shape[1]: + U, s, Vt = svd(R, full_matrices=False) + projection = Vt[:reduced_dim, :] + projected_data = (projection @ pruned_data.T).T + else: + projection = np.eye(pruned_data.shape[1]) + projected_data = pruned_data + + # Process with SSI + result = _process_ssi_block(projected_data, params, thresholds, False, 1.0) + result['Psi'] = projection + result['pruned_indices'] = idx + result['method'] = 'pcaspecx' + + except Exception as e: + result = { + 'mrep': {'imode': np.array([]), 'm0': {'lambda': np.array([]), 'shape': np.array([])}}, + 'method': 'pcaspecx_failed', + 'error': str(e) + } + + return result \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/matlab_utilities.py b/src/tokeye/eigspec/utils/matlab_utilities.py new file mode 100644 index 0000000..7efb58c --- /dev/null +++ b/src/tokeye/eigspec/utils/matlab_utilities.py @@ -0,0 +1,420 @@ +""" +Core MATLAB utility functions for eigspec package. + +This module provides Python implementations of commonly used MATLAB utility +functions from the eigspec toolbox: +- kfoldcov: K-fold cross-validation for robust covariance estimation +- logdet: Stable log-determinant computation +- srteig: Sorted eigenvalue decomposition +- zpdftmatrix: Zero-padded DFT matrix construction +- fdm1dk: FDM frequency detection with dual parameter sets + +These utilities are used throughout the eigspec analysis pipeline. +""" + +from typing import Optional, Tuple, Union, Literal, List +import numpy as np +import numpy.typing as npt +from scipy.linalg import cholesky, lu, eig, det +from scipy.stats import chi2 +import warnings + + +def kfoldcov( + X: npt.NDArray[np.floating], + k: int = 10, + criterion: Literal["logdet", "mcd", "trace", "mahalanobis"] = "mcd" +) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.floating], npt.NDArray[np.floating], Optional[npt.NDArray[np.floating]]]: + """ + K-fold cross-validation for robust covariance estimation. + + Python port of MATLAB kfoldcov.m for outlier-segment removal using + log-determinant, trace, or Mahalanobis distance metrics based on + pruned sample covariance. Subdivides data into k contiguous segments + and evaluates metrics of k pruned covariance matrices. + + Args: + X: Data matrix where each column is an observation, shape (m, n) + k: Number of segments (default 10) + criterion: Metric criterion ('logdet'/'mcd', 'trace', 'mahalanobis') + + Returns: + idx: Indices of retained observations + mu: Mean vector of pruned dataset + R: Covariance matrix of pruned dataset + md: Mahalanobis distances (optional) + """ + m, n = X.shape + + if n <= m: + raise ValueError("Must provide more data observations than data dimension") + + # Set up criterion function + if criterion in ('mcd', 'logdet'): + def FF(R, mu=None, X_seg=None): + return logdet(R, method='chol') + elif criterion == 'trace': + def FF(R, mu=None, X_seg=None): + return np.trace(R) + elif criterion == 'mahalanobis': + def FF(R, mu, X_seg): + return -_mean_mahalanobis_distance(R, mu, X_seg) + triad_argument = True + else: + warnings.warn(f"Unknown criterion '{criterion}', defaulting to 'mcd'") + def FF(R, mu=None, X_seg=None): + return logdet(R, method='chol') + + dii = n / k + ff = np.zeros(k) + min_ff = float('inf') + min_idx = None + + ii = 1 + for jj in range(k): + ii1 = int(np.round(ii)) + ii2 = int(np.round(ii + dii - 1)) + + # Create index set excluding segment jj + if jj == 0: + idx = np.arange(ii2, n) + elif jj == k - 1: + idx = np.arange(0, ii1 - 1) + else: + idx = np.concatenate([np.arange(0, ii1 - 1), np.arange(ii2, n)]) + + # Compute covariance for pruned dataset + njj = len(idx) + mu_jj = np.mean(X[:, idx], axis=1, keepdims=True) + X_jj = X[:, idx] - mu_jj + R_jj = (X_jj @ X_jj.T) / njj + + # Evaluate criterion + if criterion == 'mahalanobis': + X_seg = X[:, ii1-1:ii2] + ff[jj] = FF(R_jj, mu_jj, X_seg) + else: + ff[jj] = FF(R_jj) + + ii += dii + + # Track best result + if ff[jj] < min_ff: + min_ff = ff[jj] + min_idx = idx + + # Compute final statistics for best pruned dataset + idx = min_idx + njj = len(idx) + mu = np.mean(X[:, idx], axis=1, keepdims=True) + X_pruned = X[:, idx] - mu + R = (X_pruned @ X_pruned.T) / njj + + # Compute Mahalanobis distances if requested + md = None + if criterion == 'mahalanobis': + md = _mahalanobis_distances(R, mu, X - mu) + + return idx.astype(np.int32), mu.ravel(), R, md + + +def logdet( + A: npt.NDArray[np.floating], + method: Literal['lu', 'chol'] = 'lu' +) -> float: + """ + Stable computation of logarithm of determinant. + + Python port of MATLAB logdet.m that avoids overflow/underflow problems + when computing log(det(A)) for large matrices by using LU or Cholesky + factorization and computing the sum of log diagonal elements. + + Args: + A: Square matrix + method: Factorization method ('lu' for general, 'chol' for positive definite) + + Returns: + Log-determinant of A + """ + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A must be a square matrix") + + if method == 'chol': + # Use Cholesky factorization for positive definite matrices + try: + L = cholesky(A, lower=True) + return 2 * np.sum(np.log(np.diag(L))) + except np.linalg.LinAlgError: + # Fall back to LU if Cholesky fails + method = 'lu' + + if method == 'lu': + # Use LU factorization for general matrices + P, L, U = lu(A) + du = np.diag(U) + + # Handle potential zeros on diagonal + if np.any(du == 0): + return -np.inf + + c = det(P) * np.prod(np.sign(du)) + return np.log(np.abs(c)) + np.sum(np.log(np.abs(du))) + + raise ValueError(f"Unknown method: {method}") + + +def srteig( + A: npt.NDArray[np.floating], + sort_direction: int = 0 +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """ + Sorted eigenvalue decomposition for symmetric matrices. + + Python port of MATLAB srteig.m that returns orthonormal eigenvalue + decomposition with sorted eigenvalues. Useful for PCA of covariance matrices. + + Args: + A: Symmetric matrix + sort_direction: 0 for descending (default), 1 for ascending + + Returns: + V: Orthonormal eigenvectors (columns) + D: Sorted eigenvalues (vector) + """ + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError("A must be a square matrix") + + # Compute eigenvalue decomposition + eigenvals, eigenvecs = eig(A) + + # Sort eigenvalues and corresponding eigenvectors + if sort_direction == 1: + # Ascending order + sort_idx = np.argsort(eigenvals.real) + else: + # Descending order (default) + sort_idx = np.argsort(eigenvals.real)[::-1] + + D = eigenvals[sort_idx].real + V = eigenvecs[:, sort_idx].real + + return V, D + + +def zpdftmatrix( + M: int, + N: int +) -> Tuple[npt.NDArray[np.complex128], npt.NDArray[np.floating]]: + """ + Zero-padded DFT matrix construction. + + Python port of MATLAB zpdftmatrix.m that assembles an M-by-N matrix D + for computing zero-padded DFTs. Useful for DFT calculations with gaps + in datasets or non-standard lengths. + + Args: + M: DFT length (zero-padding if M > N) + N: Signal data length + + Returns: + D: DFT matrix, shape (M, N) + W: Angular frequencies corresponding to each row of D + """ + if M < N: + raise ValueError("Must have M >= N") + + # Angular frequencies + W = np.arange(M) * 2 * np.pi / M + + # DFT matrix using broadcasting + n_indices = np.arange(N) + D = np.exp(-1j * np.outer(W, n_indices)) + + return D, W + + +def fdm1dk( + Y: npt.NDArray[np.floating], + d1d2: Tuple[int, int], + alpha: float, + K: int, + thresh: Union[float, Tuple[float, float]] +) -> dict: + """ + FDM frequency detection with dual parameter sets. + + Python port of MATLAB fdm1dk.m that performs frequency selection based + on FDM1D routine. Makes two calls with different d-parameters and + shortlists eigenvalues that appear in both with specified threshold. + + Args: + Y: Signal data matrix + d1d2: Tuple of two d-parameters for FDM1D calls + alpha: Regularization parameter + K: Maximum number of eigenvalues to consider + thresh: Threshold value(s) for eigenvalue matching + + Returns: + Dictionary with shortlisted frequencies and errors + """ + d1, d2 = d1d2 + + if d1 == d2: + raise ValueError("d1 and d2 must be different") + + if K <= 0: + K = min(d1, d2) + + # Handle threshold parameters + if isinstance(thresh, (tuple, list)): + thresh1, thresh2 = thresh[0], thresh[1] + else: + thresh1, thresh2 = thresh, 0.95 + + if d1 < K or d2 < K: + raise ValueError(f"Both d-parameters must be >= K ({K})") + + # Call FDM1D with both parameter sets (placeholder - would need actual fdm1d implementation) + rep1 = _fdm1d_placeholder(Y, d1, alpha) + rep2 = _fdm1d_placeholder(Y, d2, alpha) + + # Extract and filter eigenvalues + e1 = rep1['eigk'] + mask1 = np.imag(e1) >= 0 + e1 = e1[mask1][:K] + err1 = rep1['errk'][mask1][:K] + + e2 = rep2['eigk'] + mask2 = np.imag(e2) >= 0 + e2 = e2[mask2][:K] + err2 = rep2['errk'][mask2][:K] + + # Find matching eigenvalues between the two sets + frq12 = [] + err12 = [] + imodes = [] + + for i, ev1 in enumerate(e1): + if err1[i] > thresh2: # Error threshold check + continue + + # Find closest match in second set + distances = np.abs(e2 - ev1) + min_idx = np.argmin(distances) + + if distances[min_idx] < (1 - thresh1) and err2[min_idx] > thresh2: + # Match found within threshold + frq12.append([ev1, e2[min_idx]]) + err12.append([err1[i], err2[min_idx]]) + imodes.append(len(frq12) - 1) + + return { + 'frq12': np.array(frq12) if frq12 else np.array([]).reshape(0, 2), + 'err12': np.array(err12) if err12 else np.array([]).reshape(0, 2), + 'imodes': np.array(imodes, dtype=np.int32) + } + + +# Helper functions + +def _mean_mahalanobis_distance( + R: npt.NDArray[np.floating], + mu: npt.NDArray[np.floating], + X: npt.NDArray[np.floating] +) -> float: + """Compute negative mean Mahalanobis distance.""" + distances = _mahalanobis_distances(R, mu, X) + return -np.mean(distances) + + +def _mahalanobis_distances( + R: npt.NDArray[np.floating], + mu: npt.NDArray[np.floating], + X: npt.NDArray[np.floating] +) -> npt.NDArray[np.floating]: + """Compute Mahalanobis distances for all columns of X.""" + m, n = X.shape + + if R.shape != (m, m): + raise ValueError("R and X are size incompatible") + + # Compute inverse of R + try: + iR = np.linalg.solve(R, np.eye(m)) + except np.linalg.LinAlgError: + iR = np.linalg.pinv(R) + + # Compute distances for all points + X_centered = X - mu + distances = np.sqrt(np.sum(X_centered * (iR @ X_centered), axis=0)) + + return distances + + +def _fdm1d_placeholder(Y: npt.NDArray[np.floating], d: int, alpha: float) -> dict: + """ + Placeholder for FDM1D function (would need full implementation). + + This is a simplified placeholder that returns mock results with the + expected structure. A full implementation would require the complete + FDM1D algorithm. + """ + N = Y.shape[0] + + # Generate mock eigenvalues (would be computed by actual FDM1D) + eigk = np.random.complex128(d) + eigk = eigk[np.argsort(np.abs(eigk))[::-1]] # Sort by magnitude + + # Mock error estimates + errk = np.random.uniform(0.8, 0.99, d) + + return { + 'eigk': eigk, + 'errk': errk + } + + +# Additional utility functions + +def periodic_smooth_1d( + data: Optional[npt.NDArray[np.floating]], + coordinates: npt.NDArray[np.floating], + period: float, + smoothing_param: float, + K: int +) -> Tuple[callable, callable, float]: + """ + Periodic 1D smoothing (simplified placeholder). + + This would be a full implementation of periodic_smooth_1d.m for + smoothing periodic data. Currently returns a placeholder. + """ + def smoothed_real(x): + return np.zeros_like(x) + + def smoothed_imag(x): + return np.zeros_like(x) + + return smoothed_real, smoothed_imag, abs(smoothing_param) + + +def weighted_rms( + data: npt.NDArray[np.floating], + weights: Optional[npt.NDArray[np.floating]] = None +) -> float: + """ + Compute weighted RMS value. + + Args: + data: Input data array + weights: Optional weight array (uniform if None) + + Returns: + Weighted RMS value + """ + if weights is None: + return np.sqrt(np.mean(data**2)) + else: + if weights.shape != data.shape: + raise ValueError("Data and weights must have same shape") + return np.sqrt(np.sum(weights * data**2) / np.sum(weights)) \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/modal_analysis.py b/src/tokeye/eigspec/utils/modal_analysis.py new file mode 100644 index 0000000..699d640 --- /dev/null +++ b/src/tokeye/eigspec/utils/modal_analysis.py @@ -0,0 +1,1053 @@ +""" +Modal analysis functions for eigspec package. + +This module provides modal analysis functions for system identification and mode shape analysis: +- Modal frequency and shape vector extraction +- Order-MAC modal matching and validation +- Shape vector estimation from frequency analysis + +Based on the MATLAB eigspec toolbox modal analysis functions: +- ac2modelist.m - Extract modal frequencies and shapes from state-space models +- order_mac() in rndspecx.m - Order-MAC modal matching procedure +- shapes_from_fshortlistx() in rndspecx.m - Shape estimation from frequency shortlist +- collect_rep_data.m - Harvest modal features from analysis reports +- mnfit_ptref.m - Point-reference mode shape fitting +- mnfit_clus_medoid.m - Cluster-based mode shape fitting +""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Dict, Union, TYPE_CHECKING, Literal + +import numpy as np +from numpy.typing import NDArray +import numpy.typing as npt + +# Import will be done with TYPE_CHECKING to avoid circular import +if TYPE_CHECKING: + from ..analysis.random_projection import RandomProjectionSpectralAnalysisResult + +if TYPE_CHECKING: + from .clustering import ClusteringResult + +@dataclass +class ModalList: + """ + Container for modal frequencies and shape vectors. + + Attributes: + eigenvalues: Complex eigenvalues representing modal frequencies + shape: Complex modal shape vectors, shape (n_outputs, n_modes) + """ + eigenvalues: NDArray[np.complexfloating] + shape: NDArray[np.complexfloating] + + # Add property aliases for backward compatibility with tests + @property + def lambda_vals(self) -> NDArray[np.complexfloating]: + """Alias for eigenvalues for test compatibility.""" + return self.eigenvalues + +@dataclass +class ModalShortlist: + """ + Container for modal analysis shortlist results. + + Attributes: + input_modes: Original modal list + mac: Modal Assurance Criterion matrix, shape (n_modes1, n_modes2) + distance: Distance matrix, shape (n_modes1, n_modes2) + mode_indices: Selected mode indices + """ + input_modes: ModalList + mac: NDArray[np.floating] + distance: NDArray[np.floating] + mode_indices: List[int] + + # Add property aliases for backward compatibility with tests + @property + def imode(self) -> List[int]: + """Alias for mode_indices for test compatibility.""" + return self.mode_indices + + @property + def m0(self) -> ModalList: + """Alias for input_modes for test compatibility.""" + return self.input_modes + + @property + def dst(self) -> NDArray[np.floating]: + """Alias for distance matrix for test compatibility.""" + return self.distance + +@dataclass +class ShapeEstimates: + """ + Container for estimated shape vectors. + + Attributes: + shape_estimates: Estimated shape matrix, shape (n_outputs, n_modes) - complex + shape_estimates_rms: RMS values for each shape, shape (n_modes,) + """ + shape_estimates: NDArray[np.complexfloating] + shape_estimates_rms: NDArray[np.floating] + + # Add property aliases for backward compatibility with tests + @property + def dhat(self) -> NDArray[np.complexfloating]: + """Alias for shape_estimates for test compatibility.""" + return self.shape_estimates + + @property + def drms(self) -> NDArray[np.floating]: + """Alias for shape_estimates_rms for test compatibility.""" + return self.shape_estimates_rms + +@dataclass +class ModalFittingResult: + """Result from modal fitting analysis. + + Attributes: + mode_numbers: Array of (m,n) mode number pairs + mac_values: MAC values for each (m,n) pair + best_fit_indices: Indices sorted by decreasing MAC value + reference_shape: Reference shape vector used for fitting + spatial_coordinates: Spatial coordinate arrays (Yxy) + query_point: Query [time, frequency] point + closest_point: Actual [time, frequency] of closest match + distance: Distance between query and closest point + """ + mode_numbers: npt.NDArray[np.int32] # (N, 2) array of (m,n) pairs + mac_values: npt.NDArray[np.floating] + best_fit_indices: npt.NDArray[np.int32] + reference_shape: npt.NDArray[np.complexfloating] + spatial_coordinates: npt.NDArray[np.floating] + query_point: npt.NDArray[np.floating] # [time_ms, freq_kHz] + closest_point: npt.NDArray[np.floating] # [time_ms, freq_kHz] + distance: float + + +@dataclass +class PrototypeExtraction: + """Container for extracted prototype modes. + + Attributes: + modes: List of individual mode information + spatial_coordinates: Spatial coordinate arrays (Yxy) + sampling_time: Sampling time in seconds + block_settings: Block analysis parameters [bss] + projection_settings: Random projection parameters [rfpn] + thresholds: Analysis thresholds + """ + modes: List[Dict[str, Union[float, npt.NDArray]]] + spatial_coordinates: npt.NDArray[np.floating] + sampling_time: float + block_settings: npt.NDArray[np.int32] + projection_settings: npt.NDArray[np.floating] + thresholds: npt.NDArray[np.floating] + + +def extract_modal_parameters( + state_matrix: NDArray[np.number], + output_matrix: NDArray[np.number], + transformation_matrix: Optional[NDArray[np.number]] = None, + imag_threshold: float = 1e-14 +) -> ModalList: + """ + Returns modal frequencies and shape-vectors for the matrix pair (A,C). + + Args: + state_matrix: State transition matrix A, shape (n_states, n_states) + output_matrix: Output matrix C, shape (n_outputs, n_states) + transformation_matrix: Optional transformation matrix U, shape (n_outputs, n_outputs) + imag_threshold: Threshold for considering imaginary part zero + + Returns: + ModalList containing eigenvalues and shape vectors + + Example: + >>> a = np.array([[0, -1], [1, 0]]) # Rotation matrix + >>> c = np.eye(2) + >>> modes = ac2modelist(a, c) + >>> np.allclose(modes.lambda_vals, [1j, -1j]) + True + """ + if state_matrix.ndim != 2 or state_matrix.shape[0] != state_matrix.shape[1]: + raise ValueError(f"A must be square, got shape {state_matrix.shape}") + if output_matrix.ndim != 2 or output_matrix.shape[1] != state_matrix.shape[0]: + raise ValueError(f"C must have shape (n_outputs, n_states), got {output_matrix.shape}") + if transformation_matrix is not None and (transformation_matrix.ndim != 2 or transformation_matrix.shape[1] != output_matrix.shape[0]): + raise ValueError(f"U must have shape (n_outputs, n_outputs), got {transformation_matrix.shape}") + + # Compute eigendecomposition + eigenvals, eigenvecs = np.linalg.eig(state_matrix) + + # Select modes with non-negative imaginary part + positive_imaginary_mask = np.imag(eigenvals) >= 0 + rr = np.where(positive_imaginary_mask)[0] + + # Extract relevant eigenvalues and shapes, ensure complex type + eigenvalues = eigenvals[rr].astype(complex) + shape = output_matrix[:, rr].astype(complex) + + # Handle real modes + real_modes = np.abs(np.imag(eigenvalues)) <= imag_threshold + eigenvalues[real_modes] = np.real(eigenvalues[real_modes]) + shape[:, real_modes] = np.real(shape[:, real_modes]) + + return ModalList(eigenvalues=eigenvalues, shape=shape) + +def order_mac( + output_matrix_1: NDArray[np.number], + state_matrix_1: NDArray[np.number], + output_matrix_2: NDArray[np.number], + state_matrix_2: NDArray[np.number], + thresh: Tuple[float, float] +) -> ModalShortlist: + """ + Return a modal shortlist using the order-MAC procedure. + + Args: + output_matrix_1: First output matrix C1, shape (n_outputs, n_states1) + state_matrix_1: First state matrix A1, shape (n_states1, n_states1) + output_matrix_2: Second output matrix C2, shape (n_outputs, n_states2) + state_matrix_2: Second state matrix A2, shape (n_states2, n_states2) + thresh: (MAC threshold, DST threshold) for mode matching + + Returns: + ModalShortlist containing analysis results + + Example: + >>> a1 = np.array([[0, -1], [1, 0]]) + >>> a2 = np.array([[0, -2], [2, 0]]) + >>> c1 = c2 = np.eye(2) + >>> result = order_mac(c1, a1, c2, a2, (0.9, 0.9)) + """ + mac_thresh, dst_thresh = thresh + + # Create lists of modal shape vectors + input_modes_1 = extract_modal_parameters(state_matrix_1, output_matrix_1) + input_modes_2 = extract_modal_parameters(state_matrix_2, output_matrix_2) + + n_modes = len(input_modes_1.eigenvalues) + + # Compute MAC values using vectorized operations + # Mode shapes should be (n_outputs, n_modes) - use shape directly + mode_shapes_1 = input_modes_1.shape # Shape (n_outputs, n_modes1) + mode_shapes_2 = input_modes_2.shape # Shape (n_outputs, n_modes2) + + # MAC = |v'w|^2 / (|v|^2 * |w|^2) + # Cross-correlation matrix: (n_modes1, n_modes2) + mac_values = mode_shapes_1.conj().T @ mode_shapes_2 # (n_modes1, n_modes2) + + # Compute norms for normalization + mode_shapes_1_norms = np.sum(np.abs(mode_shapes_1)**2, axis=0) # (n_modes1,) + mode_shapes_2_norms = np.sum(np.abs(mode_shapes_2)**2, axis=0) # (n_modes2,) + + # Broadcast norms for division: (n_modes1, 1) * (1, n_modes2) + norm_matrix = mode_shapes_1_norms[:, np.newaxis] * mode_shapes_2_norms[np.newaxis, :] + + # Handle zero norms + norm_matrix = np.where(norm_matrix == 0, 1, norm_matrix) + mac = np.abs(mac_values)**2 / norm_matrix + + # Compute distance metric + eigenvalues_1 = input_modes_1.eigenvalues[:, np.newaxis] # Shape (n_modes1, 1) + eigenvalues_2 = input_modes_2.eigenvalues[np.newaxis, :] # Shape (1, n_modes2) + distance = np.maximum(1 - np.abs(eigenvalues_1 - eigenvalues_2) / np.abs(eigenvalues_1), 0) + + # Find matching modes + max_mac_values = np.max(mac, axis=1) + i_max_mac_values = np.argmax(mac, axis=1) + max_distance_values = np.max(distance, axis=1) + i_max_distance_values = np.argmax(distance, axis=1) + + # Select modes meeting criteria + mode_indices = [i for i in range(n_modes) + if (max_mac_values[i] >= mac_thresh and + max_distance_values[i] >= dst_thresh and + i_max_mac_values[i] == i_max_distance_values[i] and + np.imag(input_modes_1.eigenvalues[i]) > 0)] + + return ModalShortlist(input_modes=input_modes_1, mac=mac, distance=distance, mode_indices=mode_indices) + +def shapes_from_freq( + signal: NDArray[np.number], + modal_shortlist: ModalShortlist +) -> Optional[ShapeEstimates]: + """ + Estimate shape vectors based on the frequency-shortlist. + + Args: + signal: Block of data, shape (n_samples, n_outputs) + modal_shortlist: Modal shortlist from order_mac + + Returns: + ShapeEstimates containing estimated shapes and RMS values, + or None if no modes were found + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*t)[:, np.newaxis] + >>> modal_shortlist = ... # Modal shortlist with one mode + >>> shapes = shapes_from_freq(y, modal_shortlist) + """ + if not modal_shortlist.mode_indices: + return None + + n_samples, n_outputs = signal.shape + n_modes = len(modal_shortlist.mode_indices) + + # Build time basis matrix efficiently + time_basis_matrix = np.arange(n_samples)[:, np.newaxis] # Shape (n, 1) + mode_angles = np.angle(modal_shortlist.input_modes.eigenvalues[modal_shortlist.mode_indices]).astype(complex) # Shape (n_modes,) + + # Create time basis for all modes at once + # Shape: (n, 2*n_modes) + time_basis_matrix = np.column_stack([ + np.cos(time_basis_matrix * mode_angles.reshape(1, -1)), + np.sin(time_basis_matrix * mode_angles.reshape(1, -1)) + ]).reshape(n_samples, -1) + + # Estimate shapes using least squares + coeffs = np.linalg.lstsq(time_basis_matrix, signal, rcond=None)[0].T + + # Convert cos/sin coefficients to complex mode shapes + # coeffs shape: (n_outputs, 2*n_modes) -> (n_outputs, n_modes, 2) + coeffs_reshaped = coeffs.reshape(n_outputs, n_modes, 2) + + # Combine cos and sin coefficients into complex form: cos_coeff + j*sin_coeff + shape_estimates = coeffs_reshaped[:, :, 0] + 1j * coeffs_reshaped[:, :, 1] + + # Compute RMS values efficiently + shape_estimates_rms = np.sqrt(np.sum(coeffs_reshaped**2, axis=(0, 2)) / n_outputs) + + return ShapeEstimates(shape_estimates=shape_estimates, shape_estimates_rms=shape_estimates_rms) + + +def extract_prototypes( + analysis_result: "RandomProjectionSpectralAnalysisResult", + query_points: npt.NDArray[np.floating], + spatial_coordinates: npt.NDArray[np.floating], + warning_distance: float = 5.0 +) -> PrototypeExtraction: + """Extract prototype mode shapes from spectral analysis results. + + Python equivalent of extract_ptrefs.m + + Finds the reference points from analysis results by locating the closest + matches to provided time-frequency query points. + + Args: + analysis_result: Random projection spectral analysis results + query_points: Query points as (N, 2) array of [time_ms, freq_kHz] + spatial_coordinates: Spatial coordinate arrays (typically Yxy) + warning_distance: Warn if distance to query point exceeds this (euclidean in ms,kHz) + + Returns: + PrototypeExtraction containing extracted mode information + + Raises: + ValueError: If inputs are malformed or analysis_result is invalid + """ + if not hasattr(analysis_result, 'block_results') or len(analysis_result.block_results) == 0: + raise ValueError("Invalid analysis_result: no block results available") + + query_points = np.asarray(query_points) + if query_points.ndim == 1: + query_points = query_points.reshape(1, -1) + + if query_points.shape[1] != 2: + raise ValueError("query_points must have shape (N, 2) for [time_ms, freq_kHz]") + + num_refs = query_points.shape[0] + + # Extract time-frequency data from all blocks + tf_data = [] + shape_vectors = [] + rms_values = [] + + # Get sampling rate from analysis result + sampling_rate = 1.0 # Default fallback + if (hasattr(analysis_result, 'block_results') and len(analysis_result.block_results) > 0 and + hasattr(analysis_result.block_results[0], 'time_step') and + analysis_result.block_results[0].time_step > 0): + sampling_rate = 1.0 / analysis_result.block_results[0].time_step + + for block_result in analysis_result.block_results: + centre_time_ms = block_result.centre_time * 1000 # Convert to ms + + # Get modal analysis results from the block - it's a ModalShortlist + modal_shortlist = block_result.reduced_dimension_matrix + if modal_shortlist is not None and hasattr(modal_shortlist, 'input_modes'): + # Access eigenvalues and mode shapes from the ModalShortlist + modal_list = getattr(modal_shortlist, 'input_modes', None) + if modal_list is not None and hasattr(modal_list, 'eigenvalues') and len(modal_list.eigenvalues) > 0: + for idx, eigenval in enumerate(modal_list.eigenvalues): + # Convert eigenvalue to frequency and damping + real_part = np.real(eigenval) + imag_part = np.imag(eigenval) + + # Frequency from imaginary part - include sampling rate for discrete-time eigenvalues + if imag_part > 0: # Only positive frequencies + freq_hz = imag_part * sampling_rate / (2 * np.pi) + freq_khz = freq_hz / 1e3 # Convert to kHz + + # Damping ratio from real part + damping = -real_part / abs(eigenval) if abs(eigenval) > 0 else 0.0 + + tf_data.append([centre_time_ms, damping, freq_khz]) + + # Get shape vector if available + modal_shape = getattr(modal_list, 'shape', None) + if (modal_shape is not None and modal_shape.size > 0 and + idx < modal_shape.shape[1]): + shape_vectors.append(modal_shape[:, idx]) + + # RMS from shape estimates if available + shape_estimates = block_result.reduced_dimension_array + if (shape_estimates is not None and + hasattr(shape_estimates, 'shape_estimates_rms')): + shape_rms = getattr(shape_estimates, 'shape_estimates_rms', None) + if shape_rms is not None and idx < len(shape_rms): + rms_values.append(shape_rms[idx]) + else: + rms_values.append(np.linalg.norm(modal_shape[:, idx])) + else: + rms_values.append(np.linalg.norm(modal_shape[:, idx])) + else: + # Create dummy shape vector + shape_vectors.append(np.zeros(len(spatial_coordinates), dtype=complex)) + rms_values.append(0.0) + + if len(tf_data) == 0: + raise ValueError("No mode data found in analysis results") + + tf_data = np.array(tf_data) + shape_vectors = np.column_stack(shape_vectors) if shape_vectors else np.array([]).reshape(len(spatial_coordinates), 0) + rms_values = np.array(rms_values) + + modes = [] + + for i in range(num_refs): + query_time, query_freq = query_points[i] + + # Find closest point in (time, frequency) space + distances_sq = ( + (tf_data[:, 0] - query_time)**2 + + (tf_data[:, 2] - query_freq)**2 + ) + + closest_idx = np.argmin(distances_sq) + distance = np.sqrt(distances_sq[closest_idx]) + + closest_time = tf_data[closest_idx, 0] + closest_freq = tf_data[closest_idx, 2] + + print(f"Query {i+1}: closest to [time,freq]=[{query_time:.6f},{query_freq:.6f}] " + f"found at index {closest_idx+1}/{len(tf_data)} " + f"[time,freq]=[{closest_time:.6f},{closest_freq:.6f}] (ms,kHz)") + + if distance > warning_distance: + print(f"(Warning: distance to query point is large, d={distance:.3f})") + + mode_info = { + 'query_tifr': np.array([query_time, query_freq]), + 'shapevector': shape_vectors[:, closest_idx].copy() if shape_vectors.size > 0 else np.array([]), + 'rms': float(rms_values[closest_idx]) if len(rms_values) > closest_idx else 0.0, + 'frequency': float(tf_data[closest_idx, 2] * 1e3), # Convert back to Hz + 'radius': float(tf_data[closest_idx, 1]), + 'centre_time': float(tf_data[closest_idx, 0] / 1e3), # Convert to seconds + 'closest_point': np.array([closest_time, closest_freq]), + 'distance': float(distance) + } + + modes.append(mode_info) + + return PrototypeExtraction( + modes=modes, + spatial_coordinates=spatial_coordinates, + sampling_time=analysis_result.block_results[0].time_step if analysis_result.block_results else 1.0, + block_settings=np.array(analysis_result.block_parameters, dtype=np.int32), + projection_settings=np.array(analysis_result.reduced_dimension, dtype=np.floating), + thresholds=np.array(analysis_result.threshold_parameters, dtype=np.floating) + ) + + +def modal_fitting_ptref( + analysis_result: "RandomProjectionSpectralAnalysisResult", + query_point: npt.NDArray[np.floating], + spatial_coordinates: npt.NDArray[np.floating], + m_max: int = 8, + n_max: int = 5, + num_display: int = 10 +) -> ModalFittingResult: + """Fit modal harmonics to reference prototype from spectral analysis. + + Python equivalent of mnfit_ptref.m + + Finds the closest reference point to the query [time, frequency] and + computes MAC values for a range of (m,n) mode numbers, displaying + the best-fitting modes. + + Args: + analysis_result: Random projection spectral analysis results + query_point: Query point as [time_ms, freq_kHz] + spatial_coordinates: Spatial coordinates (N, 2) array, typically [Y, X] or [R, Phi] + m_max: Maximum m mode number (toroidal/azimuthal) + n_max: Maximum n mode number (poloidal/radial) + num_display: Number of best fits to display/return + + Returns: + ModalFittingResult with mode fitting analysis + + Raises: + ValueError: If inputs are malformed + """ + query_point = np.asarray(query_point) + if query_point.shape != (2,): + raise ValueError("query_point must be [time_ms, freq_kHz]") + + # Extract prototype at query point + prototype = extract_prototypes( + analysis_result, query_point.reshape(1, 2), spatial_coordinates + ) + + if len(prototype.modes) == 0: + raise ValueError("No prototype modes extracted") + + mode_data = prototype.modes[0] + reference_shape = mode_data['shapevector'] + closest_point = mode_data['closest_point'] + distance = mode_data['distance'] + + # Ensure reference_shape is a proper complex numpy array + reference_shape = np.asarray(reference_shape, dtype=np.complexfloating) + # Ensure closest_point is a proper array + closest_point = np.asarray(closest_point, dtype=np.floating) + + print(f"Reference mode for fitting:") + print(f" Query: [time,freq]=[{query_point[0]:.6f},{query_point[1]:.6f}] (ms,kHz)") + print(f" Found: [time,freq]=[{closest_point[0]:.6f},{closest_point[1]:.6f}] (ms,kHz)") + print(f" Distance: {distance:.6f}") + + # Generate (m,n) mode number grid + m_vec = np.arange(-abs(m_max), abs(m_max) + 1) + n_vec = np.arange(-abs(n_max), abs(n_max) + 1) + m_grid, n_grid = np.meshgrid(m_vec, n_vec, indexing='ij') + mode_numbers = np.column_stack([m_grid.ravel(), n_grid.ravel()]) + + # Compute MAC values for all (m,n) combinations + mac_values = _compute_modal_mac_values( + reference_shape, mode_numbers, spatial_coordinates + ) + + # Sort by decreasing MAC value + sorted_indices = np.argsort(mac_values)[::-1] + + print(f"\nBest (m,n) mode fits:") + for i in range(min(num_display, len(sorted_indices))): + idx = sorted_indices[i] + m, n = mode_numbers[idx] + mac_val = mac_values[idx] + print(f" m,n={m:2d},{n:2d}; MAC={mac_val:.6f}") + + return ModalFittingResult( + mode_numbers=mode_numbers, + mac_values=mac_values, + best_fit_indices=sorted_indices.astype(np.int32), + reference_shape=reference_shape, + spatial_coordinates=spatial_coordinates, + query_point=np.asarray(query_point, dtype=np.floating), + closest_point=closest_point, + distance=float(distance) + ) + + +def modal_fitting_cluster_medoid( + data_matrix: npt.NDArray[np.complexfloating], + cluster_result: 'ClusteringResult', + cluster_index: int, + spatial_coordinates: npt.NDArray[np.floating], + m_max: int = 8, + n_max: int = 5, + num_display: int = 10 +) -> ModalFittingResult: + """Fit modal harmonics to cluster medoid shape vector. + + Python equivalent of mnfit_clus_medoid.m + + Takes a cluster medoid representative and computes MAC values for + a range of (m,n) mode numbers. + + Args: + data_matrix: Data matrix where each column is a feature vector (M, N) + cluster_result: ClusteringResult containing labels and medoid indices + cluster_index: Index of cluster to analyze (1-based as in MATLAB) + spatial_coordinates: Spatial coordinates (M, 2) array + m_max: Maximum m mode number + n_max: Maximum n mode number + num_display: Number of best fits to display/return + + Returns: + ModalFittingResult with mode fitting analysis + + Raises: + ValueError: If cluster_index is invalid or inputs malformed + """ + # Import here to avoid circular imports + from .clustering import ClusteringResult + + if not isinstance(cluster_result, ClusteringResult): + raise ValueError("cluster_result must be a ClusteringResult instance") + + if cluster_result.medoid_indices is None: + raise ValueError("cluster_result must contain medoid_indices") + + # Get number of clusters + num_clusters = len(cluster_result.medoid_indices) + + if cluster_index < 1 or cluster_index > num_clusters: + raise ValueError(f"cluster_index {cluster_index} out of range [1, {num_clusters}]") + + # Get medoid shape (convert from 1-based to 0-based indexing) + medoid_idx = cluster_result.medoid_indices[cluster_index - 1] + reference_shape = data_matrix[:, medoid_idx] + + print(f"Analyzing cluster #{cluster_index} medoid shape vector") + print(f" Medoid index: {medoid_idx}") + + # Generate (m,n) mode number grid + m_vec = np.arange(-abs(m_max), abs(m_max) + 1) + n_vec = np.arange(-abs(n_max), abs(n_max) + 1) + m_grid, n_grid = np.meshgrid(m_vec, n_vec, indexing='ij') + mode_numbers = np.column_stack([m_grid.ravel(), n_grid.ravel()]) + + # Compute MAC values for all (m,n) combinations + mac_values = _compute_modal_mac_values( + reference_shape, mode_numbers, spatial_coordinates + ) + + # Sort by decreasing MAC value + sorted_indices = np.argsort(mac_values)[::-1] + + print(f"\nBest (m,n) mode fits for cluster #{cluster_index} medoid:") + for i in range(min(num_display, len(sorted_indices))): + idx = sorted_indices[i] + m, n = mode_numbers[idx] + mac_val = mac_values[idx] + print(f" m,n={m:2d},{n:2d}; MAC={mac_val:.6f}") + + return ModalFittingResult( + mode_numbers=mode_numbers, + mac_values=mac_values, + best_fit_indices=sorted_indices.astype(np.int32), + reference_shape=reference_shape, + spatial_coordinates=spatial_coordinates, + query_point=np.array([0.0, 0.0], dtype=np.floating), # Not applicable for cluster medoids + closest_point=np.array([0.0, 0.0], dtype=np.floating), # Not applicable + distance=0.0 # Not applicable + ) + + +def _compute_modal_mac_values( + reference_shape: npt.NDArray[np.complexfloating], + mode_numbers: npt.NDArray[np.int32], + spatial_coordinates: npt.NDArray[np.floating] +) -> npt.NDArray[np.floating]: + """Compute MAC values between reference shape and (m,n) modal harmonics. + + Python equivalent of mns2mac function in mnfit_ptref.m + + Args: + reference_shape: Complex reference shape vector (M,) + mode_numbers: Array of (m,n) pairs (N, 2) + spatial_coordinates: Spatial coordinates (M, 2), typically [Y, X] + + Returns: + MAC values for each (m,n) pair (N,) + """ + M = len(reference_shape) + if M != spatial_coordinates.shape[0]: + raise ValueError("Dimension mismatch between reference_shape and spatial_coordinates") + + Y = spatial_coordinates[:, 0] # First coordinate (e.g., Y or R) + X = spatial_coordinates[:, 1] # Second coordinate (e.g., X or Phi) + + num_modes = mode_numbers.shape[0] + mac_values = np.zeros(num_modes) + + for i in range(num_modes): + m, n = mode_numbers[i] + + # Compute spatial harmonic: exp(i * (m*X + n*Y)) + # Split into cos and sin components for real computation + kx = m * X + n * Y + harmonic = np.cos(kx) + 1j * np.sin(kx) + + # Compute MAC value + mac_values[i] = _mac_value(reference_shape, harmonic) + + return mac_values + + +def _mac_value(v: npt.NDArray[np.complexfloating], w: npt.NDArray[np.complexfloating]) -> float: + """Compute Modal Assurance Criterion between two complex vectors. + + Python equivalent of macvalue function in mnfit_ptref.m + + Args: + v: First complex vector + w: Second complex vector + + Returns: + MAC value (0 to 1) + """ + numerator = (v.conj().T @ w) * (w.conj().T @ v) + denominator = (v.conj().T @ v) * (w.conj().T @ w) + + if np.abs(denominator) == 0: + return 0.0 + + return float(np.real(numerator / denominator)) + + +def complex_vector_scalar_fit( + a: npt.NDArray[np.complexfloating], + b: npt.NDArray[np.complexfloating] +) -> complex: + """Find scalar complex number c such that a ≈ c*b in least-squares sense. + + Python equivalent of complex_vector_scalar_fit.m + + Solves the complex least-squares problem to find the optimal complex scalar + that best fits one complex vector to another via scaling. + + Args: + a: Target complex vector (result of c*b) + b: Reference complex vector (to be scaled) + + Returns: + Complex scalar c that minimizes ||a - c*b||² + + Raises: + ValueError: If vectors have different lengths + """ + a = np.asarray(a).flatten() + b = np.asarray(b).flatten() + + if len(a) != len(b): + raise ValueError("Vectors a and b must have same length") + + if len(a) == 0: + return 0.0 + 0.0j + + # Solve for c in a ≈ c*b using least squares: c = / + b_conj_b = np.real(np.conj(b) @ b) + if b_conj_b > 0: + c = (np.conj(b) @ a) / b_conj_b + else: + c = 0.0 + 0.0j + + return complex(c) + + +def modal_mac_matrix( + mode_shapes: npt.NDArray[np.complexfloating] +) -> npt.NDArray[np.floating]: + """Compute MAC matrix between all pairs of mode shapes. + + Args: + mode_shapes: Mode shape matrix (n_dofs, n_modes) + + Returns: + MAC matrix (n_modes, n_modes) with MAC values between all mode pairs + """ + n_modes = mode_shapes.shape[1] + mac_matrix = np.zeros((n_modes, n_modes)) + + for i in range(n_modes): + for j in range(n_modes): + mac_matrix[i, j] = _mac_value(mode_shapes[:, i], mode_shapes[:, j]) + + return mac_matrix + + +def sort_modes_by_frequency( + frequencies: npt.NDArray[np.floating], + mode_shapes: Optional[npt.NDArray[np.complexfloating]] = None, + damping_ratios: Optional[npt.NDArray[np.floating]] = None +) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.floating], + Optional[npt.NDArray[np.complexfloating]], Optional[npt.NDArray[np.floating]]]: + """Sort modes by increasing frequency. + + Args: + frequencies: Natural frequencies (Hz) + mode_shapes: Mode shapes matrix (n_dofs, n_modes), optional + damping_ratios: Damping ratios, optional + + Returns: + Tuple of (sort_indices, sorted_frequencies, sorted_mode_shapes, sorted_damping) + """ + sort_indices = np.argsort(frequencies).astype(np.int32) + sorted_frequencies = frequencies[sort_indices] + + sorted_mode_shapes = None + if mode_shapes is not None: + sorted_mode_shapes = mode_shapes[:, sort_indices] + + sorted_damping = None + if damping_ratios is not None: + sorted_damping = damping_ratios[sort_indices] + + return sort_indices, sorted_frequencies, sorted_mode_shapes, sorted_damping + + +def normalize_mode_shapes( + mode_shapes: npt.NDArray[np.complexfloating], + method: Literal["unity_modal_mass", "max_displacement", "euclidean"] = "max_displacement" +) -> npt.NDArray[np.complexfloating]: + """Normalize mode shapes using different criteria. + + Args: + mode_shapes: Mode shapes matrix (n_dofs, n_modes) + method: Normalization method + - "unity_modal_mass": Normalize to unit modal mass (requires mass matrix) + - "max_displacement": Normalize to unit maximum displacement + - "euclidean": Normalize to unit Euclidean norm + + Returns: + Normalized mode shapes matrix + """ + normalized_shapes = mode_shapes.copy() + n_modes = mode_shapes.shape[1] + + for mode_idx in range(n_modes): + mode_shape = mode_shapes[:, mode_idx] + + if method == "max_displacement": + # Normalize by maximum absolute displacement + max_disp = np.max(np.abs(mode_shape)) + if max_disp > 0: + normalized_shapes[:, mode_idx] = mode_shape / max_disp + + elif method == "euclidean": + # Normalize to unit Euclidean norm + norm = np.linalg.norm(mode_shape) + if norm > 0: + normalized_shapes[:, mode_idx] = mode_shape / norm + + elif method == "unity_modal_mass": + # This would require mass matrix - placeholder for now + raise NotImplementedError("Unity modal mass normalization requires mass matrix") + + else: + raise ValueError(f"Unknown normalization method: {method}") + + return normalized_shapes + + +def mode_shape_scaling_factor( + reference_shape: npt.NDArray[np.complexfloating], + test_shape: npt.NDArray[np.complexfloating] +) -> complex: + """Compute optimal complex scaling factor between two mode shapes. + + Finds the complex scalar that best matches test_shape to reference_shape. + + Args: + reference_shape: Reference mode shape vector + test_shape: Test mode shape vector to be scaled + + Returns: + Complex scaling factor such that test_shape ≈ scaling_factor * reference_shape + """ + return complex_vector_scalar_fit(test_shape, reference_shape) + + +def modal_correlation_coefficient( + mode1: npt.NDArray[np.complexfloating], + mode2: npt.NDArray[np.complexfloating] +) -> float: + """Compute modal correlation coefficient between two mode shapes. + + This computes the absolute value of the normalized complex inner product, + which is appropriate for complex mode shapes and should give 0 for + orthogonal modes and 1 for identical (up to scaling) modes. + + Args: + mode1: First mode shape vector + mode2: Second mode shape vector + + Returns: + Correlation coefficient (0 to 1) + """ + if len(mode1) != len(mode2): + raise ValueError("Mode shapes must have same length") + + # Compute normalized complex inner product + numerator = np.abs(np.vdot(mode1, mode2)) + denominator = np.linalg.norm(mode1) * np.linalg.norm(mode2) + + if denominator == 0: + return 0.0 + + return float(numerator / denominator) + + +def complex_vector_scalar_fit( + a: npt.NDArray[np.complexfloating], + b: npt.NDArray[np.complexfloating] +) -> np.complexfloating: + """ + Find scalar complex number c such that a = c*b in least-squares sense. + + Python port of MATLAB complex_vector_scalar_fit.m that finds the optimal + complex scalar to fit one complex vector to another. + + Args: + a: Target complex vector + b: Reference complex vector + + Returns: + Complex scalar c such that a ≈ c*b + """ + a = a.ravel() + b = b.ravel() + + if len(a) != len(b): + raise ValueError("Vectors a and b must have same length") + + # Set up least squares system: [Re(b) -Im(b); Im(b) Re(b)] * [Re(c); Im(c)] = [Re(a); Im(a)] + M = np.block([ + [b.real.reshape(-1, 1), -b.imag.reshape(-1, 1)], + [b.imag.reshape(-1, 1), b.real.reshape(-1, 1)] + ]) + + Z = np.concatenate([a.real, a.imag]) + + # Solve least squares + ab = np.linalg.lstsq(M, Z, rcond=None)[0] + c = ab[0] + 1j * ab[1] + + return c + + +def shape2mn( + shape_vector: npt.NDArray[np.complexfloating], + mn_candidates: npt.NDArray[np.int32], + sensor_coordinates: npt.NDArray[np.floating] +) -> Tuple[int, float]: + """ + Find best-fit (m,n) harmonic for a complex shape vector. + + Python port of MATLAB shape2mn.m that determines the best (m,n) mode + numbers for a given complex shape vector by comparing with trial harmonics + using MAC (Modal Assurance Criterion). + + Args: + shape_vector: Complex shape vector, shape (M,) + mn_candidates: Trial (m,n) pairs, shape (N_trials, 2) + sensor_coordinates: Sensor angular positions, shape (M, 2) + + Returns: + J: Index of best-fit (m,n) pair (0-based) + E: MAC value of best fit + """ + M = len(shape_vector) + + if M != sensor_coordinates.shape[0]: + raise ValueError("Shape vector and sensor coordinates dimension mismatch") + + N_trials = mn_candidates.shape[0] + mac_values = np.zeros(N_trials) + + for jj in range(N_trials): + m, n = mn_candidates[jj] + + # Compute trial harmonic: exp(i*(m*theta + n*phi)) + kx = m * sensor_coordinates[:, 0] + n * sensor_coordinates[:, 1] + trial_shape = np.cos(kx) + 1j * np.sin(kx) + + # Compute MAC value + mac_values[jj] = _mac_value_util(shape_vector, trial_shape) + + # Find best match + J = np.argmax(mac_values) + E = mac_values[J] + + return int(J), float(E) + + +def _mac_value_util(v: npt.NDArray[np.complexfloating], w: npt.NDArray[np.complexfloating]) -> float: + """Compute MAC (Modal Assurance Criterion) value between two vectors.""" + numerator = np.abs(np.vdot(v, w))**2 + denominator = np.real(np.vdot(v, v) * np.vdot(w, w)) + + if denominator == 0: + return 0.0 + + return float(numerator / denominator) + + +def filter_modal_parameters( + modal_list: ModalList, + frequency_range: Optional[Tuple[float, float]] = None, + damping_threshold: Optional[float] = None +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.complexfloating]]: + """Filter modal parameters from ModalList within specified criteria. + + Args: + modal_list: ModalList containing identified modes + frequency_range: Optional (min_freq, max_freq) filter in Hz + damping_threshold: Optional maximum damping ratio filter + + Returns: + Tuple of (frequencies, damping_ratios, mode_shapes) + """ + if len(modal_list.eigenvalues) == 0: + return np.array([]), np.array([]), np.array([]).reshape(0, 0) + + # Extract frequencies and damping from eigenvalues + frequencies = [] + damping_ratios = [] + valid_indices = [] + + for i, eigenval in enumerate(modal_list.eigenvalues): + # For complex eigenvalue s = σ + jω, frequency = |ω|/(2π), damping = -σ/|s| + real_part = np.real(eigenval) + imag_part = np.imag(eigenval) + + frequency = abs(imag_part) / (2 * np.pi) + if abs(eigenval) > 0: + damping_ratio = -real_part / abs(eigenval) + else: + damping_ratio = 0.0 + + # Apply filters + if frequency_range is not None: + min_freq, max_freq = frequency_range + if not (min_freq <= frequency <= max_freq): + continue + + if damping_threshold is not None: + if damping_ratio > damping_threshold: + continue + + frequencies.append(frequency) + damping_ratios.append(damping_ratio) + valid_indices.append(i) + + if not frequencies: + return np.array([]), np.array([]), np.array([]).reshape(0, 0) + + # Extract corresponding mode shapes + if modal_list.shape.size > 0: + filtered_shapes = modal_list.shape[:, valid_indices] + else: + filtered_shapes = np.array([]).reshape(0, len(valid_indices)) + + return ( + np.array(frequencies), + np.array(damping_ratios), + filtered_shapes + ) \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/signal_processing.py b/src/tokeye/eigspec/utils/signal_processing.py new file mode 100644 index 0000000..4941f08 --- /dev/null +++ b/src/tokeye/eigspec/utils/signal_processing.py @@ -0,0 +1,2401 @@ +""" +Signal processing utilities for eigspec package. + +This module provides signal processing functions for spectral analysis and filtering: +- FFT spectral analysis and aggregation +- Digital filtering and windowing +- Coherence estimation and filtering +- Signal conditioning and preprocessing + +Based on the MATLAB eigspec toolbox signal processing functions: +- fftspec.m - Multi-channel FFT spectral analysis with aggregation +- fftspec1.m - Single-channel FFT spectral analysis +- fftspecwin.m - Windowed FFT spectral analysis +- zmfftspec.m - Zero-mean FFT spectral analysis +- kdftspec.m - K-fold DFT spectral analysis +- yfilt.m - Digital filtering with boundary handling +- yintegrate.m - Numerical integration of time series +- ydecimate.m - Decimation with anti-aliasing +- yresample.m - Resampling with interpolation +- yaddgauss.m - Add Gaussian noise to signals +- oddevenupdate_emc.m - Even/odd mode coherence estimation +""" + +from dataclasses import dataclass +from typing import Optional, Tuple, Union, Literal, List, Dict + +import numpy as np +from numpy.typing import NDArray +from scipy import signal +from scipy.fft import fft, fftfreq +import numpy.typing as npt + + +@dataclass +class FFTSpectralResult: + """ + Container for FFT spectral analysis results. + + Attributes: + P: Power spectral density matrix, shape (n_frequencies, n_blocks) + F: Frequency vector (normalized or Hz) + T: Time vector for block centers + nfft: FFT length used + block_size: Block size used + n_blocks: Number of blocks processed + fs: Sampling frequency (if provided) + """ + P: NDArray[np.floating] + F: NDArray[np.floating] + T: NDArray[np.floating] + nfft: int + block_size: int + n_blocks: int + fs: Optional[float] = None + + +@dataclass +class CoherenceResult: + """ + Container for coherence analysis results. + + Attributes: + coherence: Coherence estimates over time + time: Time vector + frequency: Frequency vector + phase: Phase difference (if computed) + """ + coherence: NDArray[np.floating] + time: NDArray[np.floating] + frequency: NDArray[np.floating] + phase: Optional[NDArray[np.floating]] = None + + +@dataclass +class AssessmentResult: + """Result from array assessment analysis. + + Attributes: + predictability: Deviation accounted for (DAF) by AR model per channel + participation: Normalized participation/relevance of each channel + rms_values: RMS values for each channel + median_scores: Median values across all time blocks [DAF, participation, RMS] + time_vector: Time vector for block centers + channel_names: Channel names/identifiers + block_parameters: Block size and stride used + ar_lag: AR lag parameter used (for AR-based assessment) + """ + predictability: npt.NDArray[np.floating] # (n_channels, n_blocks) + participation: npt.NDArray[np.floating] # (n_channels, n_blocks) + rms_values: npt.NDArray[np.floating] # (n_channels, n_blocks) + median_scores: npt.NDArray[np.floating] # (n_channels, 3) - [DAF, participation, RMS] + time_vector: npt.NDArray[np.floating] # (n_blocks,) + channel_names: List[str] + block_parameters: Tuple[int, int] + ar_lag: Optional[int] = None + + +@dataclass +class CorrelationAssessmentResult: + """Result from correlation-based array assessment. + + Attributes: + correlation_matrix: Cross-correlation matrix between channels + channel_names: Channel names/identifiers + block_parameters: Block size and stride used + time_vector: Time vector for analysis + """ + correlation_matrix: npt.NDArray[np.floating] # (n_channels, n_channels) + channel_names: List[str] + block_parameters: Tuple[int, int] + time_vector: npt.NDArray[np.floating] + + +def compute_zero_mean_spectrum( + time_array: NDArray[np.number], + signal_array: NDArray[np.number], + block_size: int, + nfft: Union[int, Tuple[int, int]], + reduced_dim: Optional[int] = None, + figure_num: Optional[int] = None, + freq_span: Optional[Tuple[float, float]] = None +) -> FFTSpectralResult: + """ + Basic FFT-derived power spectral density plot of multivariate time-series. + + This is a port of the MATLAB zmfftspec function that aggregates channel FFTs + and provides power spectral density estimates using block processing. + + Args: + time_array: Time vector, shape (n_samples,) + signal_array: Multi-channel data, shape (n_samples, n_channels) + block_size: Block size for processing + nfft: FFT length, or tuple (fft_length, smoothing_span) + reduced_dim: Optional dimension reduction (< n_channels) + figure_num: Figure number for plotting (ignored in Python version) + freq_span: Optional frequency range for analysis + + Returns: + FFTSpectralResult containing power spectral density and metadata + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*t)[:, np.newaxis] + >>> result = compute_zero_mean_spectrum(t, y, block_size=100, nfft=128) + >>> print(f"PSD shape: {result.P.shape}") + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + if len(time_array) != signal_array.shape[0]: + raise ValueError("time_array length must match signal_array rows") + + n_samples, n_channels = signal_array.shape + + # Handle nfft specification + if isinstance(nfft, (list, tuple)) and len(nfft) == 2: + smooth_span = nfft[1] + nfft = nfft[0] + else: + smooth_span = 1 + + if nfft < block_size: + raise ValueError("NFFT must be >= block_size") + if nfft % 2 != 0: + raise ValueError("NFFT should be an even number") + + nfft_half = nfft // 2 + + # Calculate block parameters + n_blocks = (n_samples - block_size) // block_size + 1 + dt = time_array[1] - time_array[0] if len(time_array) > 1 else 1.0 + fs = 1.0 / dt + + # Initialize output arrays + P = np.zeros((nfft_half, n_blocks)) + T = np.zeros(n_blocks) + + # Frequency vector (normalized angular frequency) + F = np.arange(nfft_half) * (2 * np.pi) / nfft + + # Process each block + for j in range(n_blocks): + start_idx = j * block_size + end_idx = start_idx + block_size + T[j] = time_array[start_idx + block_size // 2] # Block center time + + # Extract block and compute aggregated FFT + Y_block = signal_array[start_idx:end_idx, :] + P[:, j] = aggregate_fft(Y_block, nfft, reduced_dim) + + # Apply smoothing if requested + if smooth_span > 1: + # Simple moving average smoothing + kernel = np.ones(smooth_span) / smooth_span + for i in range(nfft_half): + P[i, :] = np.convolve(P[i, :], kernel, mode='same') + + return FFTSpectralResult( + P=P, F=F, T=T, nfft=nfft, + block_size=block_size, n_blocks=n_blocks, fs=fs + ) + + +def aggregate_fft( + signal_block: NDArray[np.number], + nfft: int, + reduced_dim: Optional[int] = None +) -> NDArray[np.floating]: + """ + Aggregate power spectrum from all channels in a signal block. + + Args: + signal_block: Block of multichannel data, shape (n_samples, n_channels) + nfft: FFT length + reduced_dim: Optional dimension for random projection + + Returns: + Aggregated power spectrum, shape (nfft//2,) + """ + n_samples, n_channels = signal_block.shape + nfft_half = nfft // 2 + + if reduced_dim is not None and reduced_dim < n_channels: + # Apply random projection for dimension reduction + projection_matrix = np.random.randn(reduced_dim, n_channels) + projection_matrix /= np.linalg.norm(projection_matrix, axis=1, keepdims=True) + signal_block = (projection_matrix @ signal_block.T).T + n_channels = reduced_dim + + # Compute FFT for each channel and aggregate + P = np.zeros(nfft_half) + for ch in range(n_channels): + Y_fft = fft(signal_block[:, ch], nfft) + Y_half = Y_fft[:nfft_half] + P += np.abs(Y_half) ** 2 + + return P / n_channels + + +def filter_signal( + signal_array: NDArray[np.number], + filter_type: Literal['LP', 'HP', 'BP', 'BS'], + frequency: Union[float, Tuple[float, float]], + filter_method: Literal['filtfilt', 'filter'] = 'filtfilt', + filter_order: int = 2 +) -> NDArray[np.floating]: + """ + Filter columns of signal array independently as time-series. + + Supports LP (lowpass), HP (highpass), BP (bandpass), and BS (bandstop) filters + using Butterworth design by default. + + Args: + signal_array: Input data, shape (n_samples, n_channels) + filter_type: Type of filter ('LP', 'HP', 'BP', 'BS') + frequency: Cutoff frequency (scalar) or band (tuple) in normalized units [0,1] + filter_method: 'filtfilt' for zero-phase or 'filter' for causal + filter_order: Filter order (default: 2) + + Returns: + Filtered signal array, same shape as input + + Example: + >>> y = np.random.randn(1000, 3) + >>> y_filt = filter_signal(y, 'LP', 0.1) # 10% Nyquist lowpass + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + + filter_type = filter_type.upper() + if filter_type not in ['LP', 'HP', 'BP', 'BS']: + raise ValueError(f"Invalid filter_type: {filter_type}") + + # Validate frequency specification + if filter_type in ['BP', 'BS']: + if not isinstance(frequency, (list, tuple)) or len(frequency) != 2: + raise ValueError(f"{filter_type} filter requires frequency band [f1, f2]") + if any(f <= 0 or f >= 1 for f in frequency): + raise ValueError("Frequencies must be in range (0, 1)") + else: + if isinstance(frequency, (list, tuple)): + frequency = frequency[0] + if frequency <= 0 or frequency >= 1: + raise ValueError("Frequency must be in range (0, 1)") + + # Design Butterworth filter + if filter_type == 'LP': + sos = signal.butter(filter_order, frequency, btype='low', output='sos') + elif filter_type == 'HP': + sos = signal.butter(filter_order, frequency, btype='high', output='sos') + elif filter_type == 'BP': + sos = signal.butter(filter_order, frequency, btype='band', output='sos') + elif filter_type == 'BS': + sos = signal.butter(filter_order, frequency, btype='bandstop', output='sos') + + # Apply filter to each channel + n_samples, n_channels = signal_array.shape + filtered_signal = np.zeros_like(signal_array) + + for ch in range(n_channels): + if filter_method == 'filtfilt': + filtered_signal[:, ch] = signal.sosfiltfilt(sos, signal_array[:, ch]) + else: + filtered_signal[:, ch] = signal.sosfilt(sos, signal_array[:, ch]) + + return filtered_signal + + +def create_window( + window_type: Literal['hann', 'hamming', 'rectangular', 'blackman'], + window_length: int +) -> NDArray[np.floating]: + """ + Create window function for spectral analysis. + + Args: + window_type: Type of window function + window_length: Length of window + + Returns: + Window function values, shape (window_length,) + """ + if window_type == 'hann': + return signal.windows.hann(window_length) + elif window_type == 'hamming': + return signal.windows.hamming(window_length) + elif window_type == 'rectangular': + return np.ones(window_length) + elif window_type == 'blackman': + return signal.windows.blackman(window_length) + else: + raise ValueError(f"Unknown window type: {window_type}") + + +def fftspec( + time_array: NDArray[np.number], + signal_array: NDArray[np.number], + block_size: int, + nfft: int, + reduced_dim: Optional[int] = None, + window_type: str = 'hann' +) -> FFTSpectralResult: + """ + FFT spectral analysis with windowing and optional dimension reduction. + + Args: + time_array: Time vector + signal_array: Multi-channel signal data + block_size: Block size for analysis + nfft: FFT length + reduced_dim: Optional dimension reduction + window_type: Window function type + + Returns: + FFTSpectralResult with spectral analysis results + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + + n_samples, n_channels = signal_array.shape + n_blocks = (n_samples - block_size) // block_size + 1 + nfft_half = nfft // 2 + + # Create window function + window = create_window(window_type, block_size) + window_power = np.sum(window ** 2) + + # Initialize output + P = np.zeros((nfft_half, n_blocks)) + T = np.zeros(n_blocks) + + # Frequency vector + dt = time_array[1] - time_array[0] if len(time_array) > 1 else 1.0 + F = fftfreq(nfft, dt)[:nfft_half] + + # Process each block + for j in range(n_blocks): + start_idx = j * block_size + end_idx = start_idx + block_size + T[j] = time_array[start_idx + block_size // 2] + + # Extract and window the block + block = signal_array[start_idx:end_idx, :] + windowed_block = block * window[:, np.newaxis] + + # Compute power spectrum + P[:, j] = aggregate_fft(windowed_block, nfft, reduced_dim) + + # Normalize by window power + P /= window_power + + return FFTSpectralResult( + P=P, F=F, T=T, nfft=nfft, + block_size=block_size, n_blocks=n_blocks, fs=1.0/dt + ) + + +def coherence_filter( + signal_x: NDArray[np.number], + signal_y: NDArray[np.number], + block_size: int, + nfft: int, + stride: int, + forgetting_factor: float +) -> CoherenceResult: + """ + Coherence filter with forgetting factor for two signals. + + Args: + signal_x: First signal + signal_y: Second signal + block_size: Block size for analysis + nfft: FFT length + stride: Stride between blocks + forgetting_factor: Forgetting factor (0 < beta < 1) + + Returns: + CoherenceResult with coherence estimates over time + """ + if len(signal_x) != len(signal_y): + raise ValueError("Signals must have same length") + if nfft < block_size: + raise ValueError("nfft must be >= block_size") + if nfft % 2 != 0: + raise ValueError("nfft must be even") + + n_samples = len(signal_x) + nfft_half = nfft // 2 + n_blocks = (n_samples - block_size) // stride + 1 + + # Initialize + coherence = np.zeros((nfft_half, n_blocks)) + time_vec = np.zeros(n_blocks) + frequency = np.arange(nfft_half) * np.pi / nfft_half + + # Create Hann window + window = signal.windows.hann(block_size) + + # Initialize filtered cross-spectral quantities + Pxx = np.zeros(nfft_half, dtype=complex) + Pyy = np.zeros(nfft_half, dtype=complex) + Pxy = np.zeros(nfft_half, dtype=complex) + + for j in range(n_blocks): + start_idx = j * stride + end_idx = start_idx + block_size + time_vec[j] = start_idx + block_size // 2 + + # Extract and window signals + x_block = signal_x[start_idx:end_idx] * window + y_block = signal_y[start_idx:end_idx] * window + + # Compute FFTs + X = fft(x_block, nfft)[:nfft_half] + Y = fft(y_block, nfft)[:nfft_half] + + # Update filtered quantities with forgetting factor + Pxx = forgetting_factor * Pxx + (1 - forgetting_factor) * (X * X.conj()) + Pyy = forgetting_factor * Pyy + (1 - forgetting_factor) * (Y * Y.conj()) + Pxy = forgetting_factor * Pxy + (1 - forgetting_factor) * (X * Y.conj()) + + # Compute coherence with numerical stability + denominator = np.abs(Pxx) * np.abs(Pyy) + # Avoid division by zero and ensure coherence <= 1 + coherence[:, j] = np.minimum(1.0, np.abs(Pxy) ** 2 / np.maximum(denominator, 1e-15)) + + return CoherenceResult( + coherence=coherence, time=time_vec, + frequency=frequency, phase=np.angle(Pxy) + ) + + +def ar_assessment( + time_vector: Optional[npt.NDArray[np.floating]], + signal_array: npt.NDArray[np.floating], + ar_lag: int, + block_parameters: Tuple[int, int], + median_filter_length: int = 0, + channel_names: Optional[List[str]] = None +) -> AssessmentResult: + """AR-based array assessment for detecting problematic channels. + + Python equivalent of arassess.m + + Analyzes each channel's predictability using autoregressive modeling + and participation in the overall array response. Useful for identifying + faulty sensors or channels with poor signal quality. + + Args: + time_vector: Time vector or None (uses indices if None) + signal_array: Input data matrix (n_samples, n_channels) + ar_lag: AR model lag (past samples to use for prediction) + block_parameters: (block_size, block_stride) for analysis + median_filter_length: Length of median filter preprocessing (0 = none) + channel_names: Optional channel names (generated if None) + + Returns: + AssessmentResult with channel health metrics + + Raises: + ValueError: If inputs are malformed or inconsistent + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + + n_samples, n_channels = signal_array.shape + + if time_vector is not None: + time_vector = np.asarray(time_vector) + if len(time_vector) != n_samples: + raise ValueError("Length of time_vector must match signal_array first dimension") + + if ar_lag < 1: + raise ValueError("ar_lag must be positive") + + block_size, block_stride = block_parameters + if block_size <= ar_lag: + raise ValueError("block_size must be larger than ar_lag") + + # Generate channel names if not provided + if channel_names is None: + channel_names = [f"#{i+1}" for i in range(n_channels)] + elif len(channel_names) != n_channels: + raise ValueError("Length of channel_names must match number of channels") + + # Add channel indices to names for reference + channel_names = [f"{name} (#{i+1})" for i, name in enumerate(channel_names)] + + # Normalize channels to unit RMS + Y = signal_array.copy() + channel_rms = np.sqrt(np.mean(Y**2, axis=0)) + + # Avoid division by zero + nonzero_rms = channel_rms > 0 + Y[:, nonzero_rms] = Y[:, nonzero_rms] / channel_rms[nonzero_rms] + + if not np.all(nonzero_rms): + print(f"Warning: {np.sum(~nonzero_rms)} channels have zero RMS") + + # Check for rank deficiency + try: + cov_matrix = np.cov(Y.T) + rank = np.linalg.matrix_rank(cov_matrix) + if rank < n_channels: + print(f"Warning: signal data is not full rank (rank={rank}, channels={n_channels})") + except np.linalg.LinAlgError: + print("Warning: could not compute covariance matrix rank") + + # Apply median filtering if requested + if median_filter_length >= 3: + from scipy.signal import medfilt + for ch in range(n_channels): + Y[:, ch] = medfilt(Y[:, ch], kernel_size=median_filter_length) + + # Block analysis setup + block_starts = np.arange(0, n_samples - block_size + 1, block_stride) + n_blocks = len(block_starts) + + if n_blocks == 0: + raise ValueError("No complete blocks can be formed with given parameters") + + # Initialize results + predictability = np.zeros((n_channels, n_blocks)) # Deviation Accounted For (DAF) + participation = np.zeros((n_channels, n_blocks)) # Channel participation/relevance + rms_block_values = np.zeros((n_channels, n_blocks)) # Block RMS values + + if time_vector is not None: + time_centers = np.zeros(n_blocks) + else: + time_centers = np.arange(n_blocks) + + for block_idx, block_start in enumerate(block_starts): + block_end = block_start + block_size + + if time_vector is not None: + time_centers[block_idx] = (time_vector[block_start] + time_vector[block_end-1]) / 2 + + # Extract block data + block_data = Y[block_start:block_end, :] + + # Compute AR model using multivariate approach + try: + ar_matrices, cov_innovation, cov_data = _multivariate_ar_model(block_data, ar_lag) + + # Compute predictability (Deviation Accounted For) + innovation_var = np.diag(cov_innovation) + data_var = np.diag(cov_data) + + # Avoid division by zero + valid_vars = data_var > 0 + daf = np.zeros(n_channels) + daf[valid_vars] = np.maximum(0, 1 - np.sqrt(innovation_var[valid_vars] / data_var[valid_vars])) + predictability[:, block_idx] = daf + + # Compute participation (contribution to AR model) + participation_scores = np.zeros(n_channels) + for ch in range(n_channels): + # Sum of L2 norms across all AR matrices for this channel + channel_contrib = 0.0 + for lag in range(ar_lag): + # Use default norm (L2) for vectors, not Frobenius norm + channel_contrib += np.linalg.norm(ar_matrices[lag][:, ch])**2 + participation_scores[ch] = np.sqrt(channel_contrib) + + # Normalize participation + total_participation = np.sum(participation_scores) + if total_participation > 0: + participation[:, block_idx] = n_channels * participation_scores / total_participation + else: + participation[:, block_idx] = np.ones(n_channels) # Equal participation if all zero + + except (np.linalg.LinAlgError, ValueError) as e: + print(f"Warning: AR model failed for block {block_idx}: {e}") + # Set default values for failed blocks + predictability[:, block_idx] = 0.0 + participation[:, block_idx] = 1.0 + + # Store RMS values (scaled back to original units) + block_rms = np.sqrt(np.mean(block_data**2, axis=0)) + rms_block_values[:, block_idx] = block_rms * channel_rms + + # Compute median scores across all blocks + median_scores = np.column_stack([ + np.median(predictability, axis=1), + np.median(participation, axis=1), + np.median(rms_block_values, axis=1) + ]) + + return AssessmentResult( + predictability=predictability, + participation=participation, + rms_values=rms_block_values, + median_scores=median_scores, + time_vector=time_centers, + channel_names=channel_names, + block_parameters=block_parameters, + ar_lag=ar_lag + ) + + +def correlation_assessment( + time_vector: Optional[npt.NDArray[np.floating]], + signal_array: npt.NDArray[np.floating], + block_parameters: Tuple[int, int], + channel_names: Optional[List[str]] = None +) -> CorrelationAssessmentResult: + """Simple correlation-based array assessment. + + Python equivalent of corrassess.m + + Computes cross-correlation matrix between all channels to assess + array coherence and identify outlier channels. + + Args: + time_vector: Time vector or None (uses indices if None) + signal_array: Input data matrix (n_samples, n_channels) + block_parameters: (block_size, block_stride) for analysis + channel_names: Optional channel names (generated if None) + + Returns: + CorrelationAssessmentResult with correlation matrix + + Raises: + ValueError: If inputs are malformed + """ + if signal_array.ndim != 2: + raise ValueError(f"signal_array must be 2D, got shape {signal_array.shape}") + + n_samples, n_channels = signal_array.shape + + if time_vector is not None: + time_vector = np.asarray(time_vector) + if len(time_vector) != n_samples: + raise ValueError("Length of time_vector must match signal_array first dimension") + else: + time_vector = np.arange(n_samples) + + # Generate channel names if not provided + if channel_names is None: + channel_names = [f"#{i+1}" for i in range(n_channels)] + elif len(channel_names) != n_channels: + raise ValueError("Length of channel_names must match number of channels") + + block_size, block_stride = block_parameters + + # Use middle section of data for correlation analysis + if block_size < n_samples: + start_idx = (n_samples - block_size) // 2 + end_idx = start_idx + block_size + analysis_data = signal_array[start_idx:end_idx, :] + analysis_time = time_vector[start_idx:end_idx] + else: + analysis_data = signal_array + analysis_time = time_vector + + print(f"Computing correlation matrix for {n_channels} channels using {len(analysis_data)} samples") + + # Compute correlation matrix with warning suppression for constant signals + with np.errstate(divide='ignore', invalid='ignore'): + correlation_matrix = np.corrcoef(analysis_data.T) + + # Handle NaN values (can occur with constant signals) + if np.any(np.isnan(correlation_matrix)): + print("Warning: NaN values detected in correlation matrix (replacing with zeros)") + correlation_matrix = np.nan_to_num(correlation_matrix, nan=0.0) + + return CorrelationAssessmentResult( + correlation_matrix=correlation_matrix, + channel_names=channel_names, + block_parameters=block_parameters, + time_vector=analysis_time + ) + + +def _multivariate_ar_model( + data: npt.NDArray[np.floating], + lag: int +) -> Tuple[List[npt.NDArray[np.floating]], npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """Fit multivariate AR model to data. + + Python equivalent of the subvar function in arassess.m + + Args: + data: Input data matrix (n_samples, n_channels) + lag: AR model lag order + + Returns: + Tuple of (ar_matrices, innovation_covariance, data_covariance) + - ar_matrices: List of AR coefficient matrices [A1, A2, ..., Ap] + - innovation_covariance: Covariance of prediction errors + - data_covariance: Covariance of input data + + Raises: + ValueError: If data is insufficient or ill-conditioned + """ + n_samples, n_channels = data.shape + + if n_samples <= lag: + raise ValueError(f"Need more samples than lag: {n_samples} <= {lag}") + + n_effective = n_samples - lag + n_features = n_channels * lag + + # Build regressor matrix (past observations) and response matrix (current observations) + Z = data.T # Transpose for easier indexing + regressor_matrix = np.zeros((n_features, n_effective)) + response_matrix = np.zeros((n_channels, n_effective)) + + for k in range(lag, n_samples): + # Stack past lag observations into regressor vector + past_obs = [] + for lag_idx in range(lag): + past_obs.append(Z[:, k - lag_idx - 1]) + regressor_vector = np.concatenate(past_obs) + + regressor_matrix[:, k - lag] = regressor_vector + response_matrix[:, k - lag] = data[k, :] + + # Solve normal equations: Y = H * Zp, where H are AR coefficients + try: + regressor_cov = regressor_matrix @ regressor_matrix.T / n_effective + cross_cov = response_matrix @ regressor_matrix.T / n_effective + + # Solve for AR coefficient matrix H + ar_coeff_matrix = cross_cov @ np.linalg.pinv(regressor_cov) + + except np.linalg.LinAlgError: + raise ValueError("Failed to solve AR normal equations (singular covariance matrix)") + + # Split AR coefficient matrix into per-lag matrices + ar_matrices = [] + for lag_idx in range(lag): + start_idx = lag_idx * n_channels + end_idx = (lag_idx + 1) * n_channels + ar_matrices.append(ar_coeff_matrix[:, start_idx:end_idx]) + + # Compute covariances + data_covariance = response_matrix @ response_matrix.T / n_effective + + # Prediction errors + prediction_errors = response_matrix - ar_coeff_matrix @ regressor_matrix + innovation_covariance = prediction_errors @ prediction_errors.T / n_effective + + return ar_matrices, innovation_covariance, data_covariance + + +try: + from numba import jit, prange + _NUMBA_AVAILABLE = True +except ImportError: + _NUMBA_AVAILABLE = False + + # Fallback decorator that does nothing + def jit(*args, **kwargs): + def decorator(func): + return func + return decorator + + def prange(x): + return range(x) + +try: + from joblib import Parallel, delayed + _JOBLIB_AVAILABLE = True +except ImportError: + _JOBLIB_AVAILABLE = False + + # Fallback implementations + def delayed(func): + return func + + class Parallel: + def __init__(self, *args, **kwargs): + pass + + def __call__(self, iterable): + return list(iterable) + + +@jit(nopython=True, cache=True) if _NUMBA_AVAILABLE else lambda x: x +def _fast_fft_magnitude_squared(real_part: npt.NDArray[np.floating], + imag_part: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: + """Optimized computation of FFT magnitude squared. + + Args: + real_part: Real part of FFT + imag_part: Imaginary part of FFT + + Returns: + Magnitude squared values + """ + result = np.zeros_like(real_part) + for i in prange(len(real_part)): + result[i] = real_part[i]**2 + imag_part[i]**2 + return result + + +@jit(nopython=True, cache=True) if _NUMBA_AVAILABLE else lambda x: x +def _fast_mac_computation(shape1_real: npt.NDArray[np.floating], + shape1_imag: npt.NDArray[np.floating], + shape2_real: npt.NDArray[np.floating], + shape2_imag: npt.NDArray[np.floating]) -> float: + """Optimized MAC computation for complex vectors. + + Args: + shape1_real: Real part of first shape vector + shape1_imag: Imaginary part of first shape vector + shape2_real: Real part of second shape vector + shape2_imag: Imaginary part of second shape vector + + Returns: + MAC value + """ + # Compute complex dot products + dot_12_real = 0.0 + dot_12_imag = 0.0 + norm1_sq = 0.0 + norm2_sq = 0.0 + + for i in range(len(shape1_real)): + # Conjugate of shape1 dot shape2 + dot_12_real += shape1_real[i] * shape2_real[i] + shape1_imag[i] * shape2_imag[i] + dot_12_imag += shape1_real[i] * shape2_imag[i] - shape1_imag[i] * shape2_real[i] + + # Norms squared + norm1_sq += shape1_real[i]**2 + shape1_imag[i]**2 + norm2_sq += shape2_real[i]**2 + shape2_imag[i]**2 + + # MAC = |dot_12|^2 / (norm1_sq * norm2_sq) + dot_12_mag_sq = dot_12_real**2 + dot_12_imag**2 + + if norm1_sq * norm2_sq == 0.0: + return 0.0 + + return dot_12_mag_sq / (norm1_sq * norm2_sq) + + +@jit(nopython=True, cache=True) if _NUMBA_AVAILABLE else lambda x: x +def _fast_correlation_matrix(data: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: + """Optimized correlation matrix computation. + + Args: + data: Data matrix (n_samples, n_channels) + + Returns: + Correlation matrix (n_channels, n_channels) + """ + n_samples, n_channels = data.shape + + # Compute means + means = np.zeros(n_channels) + for j in range(n_channels): + for i in range(n_samples): + means[j] += data[i, j] + means[j] /= n_samples + + # Compute standard deviations + stds = np.zeros(n_channels) + for j in range(n_channels): + for i in range(n_samples): + diff = data[i, j] - means[j] + stds[j] += diff * diff + stds[j] = np.sqrt(stds[j] / (n_samples - 1)) + + # Compute correlation matrix + corr_matrix = np.zeros((n_channels, n_channels)) + + for i in range(n_channels): + for j in range(i, n_channels): + if stds[i] == 0.0 or stds[j] == 0.0: + corr_matrix[i, j] = 0.0 + else: + covariance = 0.0 + for k in range(n_samples): + covariance += (data[k, i] - means[i]) * (data[k, j] - means[j]) + covariance /= (n_samples - 1) + + correlation = covariance / (stds[i] * stds[j]) + corr_matrix[i, j] = correlation + corr_matrix[j, i] = correlation # Symmetric matrix + + return corr_matrix + + +def parallel_block_processing( + data: npt.NDArray[np.floating], + block_function: callable, + block_parameters: Tuple[int, int], + n_jobs: int = -1, + **kwargs +) -> List: + """Process data blocks in parallel. + + Args: + data: Input data matrix (n_samples, n_channels) + block_function: Function to apply to each block + block_parameters: (block_size, block_stride) + n_jobs: Number of parallel jobs (-1 for all cores) + **kwargs: Additional arguments for block_function + + Returns: + List of results from each block + """ + if not _JOBLIB_AVAILABLE: + print("Warning: joblib not available, falling back to sequential processing") + + n_samples = data.shape[0] + block_size, block_stride = block_parameters + + block_starts = np.arange(0, n_samples - block_size + 1, block_stride) + + def process_block(start_idx): + end_idx = start_idx + block_size + block_data = data[start_idx:end_idx] + return block_function(block_data, **kwargs) + + if _JOBLIB_AVAILABLE and n_jobs != 1: + # Parallel processing + results = Parallel(n_jobs=n_jobs)( + delayed(process_block)(start_idx) for start_idx in block_starts + ) + else: + # Sequential processing + results = [process_block(start_idx) for start_idx in block_starts] + + return results + + +def optimized_spectral_analysis( + signal_data: npt.NDArray[np.floating], + fft_size: int = 2048, + overlap: float = 0.5, + window: str = "hann", + n_jobs: int = -1 +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """Optimized parallel spectral analysis. + + Args: + signal_data: Input signals (n_samples, n_channels) + fft_size: FFT size for analysis + overlap: Overlap fraction between segments + window: Window function name + n_jobs: Number of parallel jobs + + Returns: + Tuple of (frequencies, times, power_spectral_density) + """ + n_samples, n_channels = signal_data.shape + hop_size = int(fft_size * (1 - overlap)) + + # Generate window + if window == "hann": + win = np.hanning(fft_size) + elif window == "hamming": + win = np.hamming(fft_size) + elif window == "blackman": + win = np.blackman(fft_size) + else: + win = np.ones(fft_size) # Rectangular + + # Normalize window + win_norm = np.sum(win**2) + + # Block processing function + def compute_segment_psd(data_segment): + if len(data_segment) < fft_size: + # Pad with zeros + padded = np.zeros((fft_size, n_channels)) + padded[:len(data_segment)] = data_segment + data_segment = padded + + # Apply window and compute FFT for all channels + windowed = data_segment * win[:, np.newaxis] + fft_result = fft(windowed, axis=0) + + # Compute power spectral density + if _NUMBA_AVAILABLE: + psd = np.zeros((fft_size, n_channels)) + for ch in range(n_channels): + psd[:, ch] = _fast_fft_magnitude_squared( + fft_result[:, ch].real, fft_result[:, ch].imag + ) + else: + psd = np.abs(fft_result)**2 + + # Normalize + psd = psd / (win_norm * n_samples) + + return psd + + # Parallel processing of segments + segment_starts = np.arange(0, n_samples - fft_size + 1, hop_size) + + if _JOBLIB_AVAILABLE and n_jobs != 1: + psd_segments = Parallel(n_jobs=n_jobs)( + delayed(compute_segment_psd)(signal_data[start:start+fft_size]) + for start in segment_starts + ) + else: + psd_segments = [ + compute_segment_psd(signal_data[start:start+fft_size]) + for start in segment_starts + ] + + # Combine results + psd_array = np.stack(psd_segments, axis=2) # (freq, channels, time) + + # Generate frequency and time axes + frequencies = fftfreq(fft_size, d=1.0)[:fft_size//2] # Positive frequencies only + times = segment_starts / n_samples # Normalized time + + # Return only positive frequencies + psd_positive = psd_array[:fft_size//2, :, :] + + return frequencies, times, psd_positive + + +def optimized_ar_assessment( + signal_data: npt.NDArray[np.floating], + ar_lag: int, + block_parameters: Tuple[int, int], + n_jobs: int = -1 +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """Optimized parallel AR-based assessment. + + Args: + signal_data: Input signals (n_samples, n_channels) + ar_lag: AR model lag order + block_parameters: (block_size, block_stride) + n_jobs: Number of parallel jobs + + Returns: + Tuple of (predictability_matrix, participation_matrix) + """ + def ar_block_analysis(block_data): + """Analyze single block for AR metrics.""" + try: + ar_matrices, cov_innovation, cov_data = _multivariate_ar_model(block_data, ar_lag) + + # Compute predictability + innovation_var = np.diag(cov_innovation) + data_var = np.diag(cov_data) + + valid_vars = data_var > 0 + daf = np.zeros(len(data_var)) + daf[valid_vars] = np.maximum(0, 1 - np.sqrt(innovation_var[valid_vars] / data_var[valid_vars])) + + # Compute participation + n_channels = block_data.shape[1] + participation_scores = np.zeros(n_channels) + + for ch in range(n_channels): + channel_contrib = 0.0 + for lag_idx in range(ar_lag): + channel_contrib += np.linalg.norm(ar_matrices[lag_idx][:, ch])**2 + participation_scores[ch] = np.sqrt(channel_contrib) + + # Normalize participation + total_participation = np.sum(participation_scores) + if total_participation > 0: + participation_scores = n_channels * participation_scores / total_participation + else: + participation_scores = np.ones(n_channels) + + return daf, participation_scores + + except (np.linalg.LinAlgError, ValueError): + # Return zeros for failed blocks + n_channels = block_data.shape[1] + return np.zeros(n_channels), np.ones(n_channels) + + # Use parallel block processing + results = parallel_block_processing( + signal_data, + ar_block_analysis, + block_parameters, + n_jobs=n_jobs + ) + + # Combine results + if results: + predictability_list, participation_list = zip(*results) + predictability_matrix = np.column_stack(predictability_list) + participation_matrix = np.column_stack(participation_list) + else: + n_channels = signal_data.shape[1] + predictability_matrix = np.array([]).reshape(n_channels, 0) + participation_matrix = np.array([]).reshape(n_channels, 0) + + return predictability_matrix, participation_matrix + + +def batch_mac_computation( + reference_shapes: npt.NDArray[np.complexfloating], + test_shapes: npt.NDArray[np.complexfloating], + n_jobs: int = -1 +) -> npt.NDArray[np.floating]: + """Compute MAC values between reference and test shapes in parallel. + + Args: + reference_shapes: Reference mode shapes (n_dofs, n_ref_modes) + test_shapes: Test mode shapes (n_dofs, n_test_modes) + n_jobs: Number of parallel jobs + + Returns: + MAC matrix (n_ref_modes, n_test_modes) + """ + n_ref_modes = reference_shapes.shape[1] + n_test_modes = test_shapes.shape[1] + + def compute_mac_row(ref_idx): + """Compute MAC values for one reference mode against all test modes.""" + ref_shape = reference_shapes[:, ref_idx] + mac_row = np.zeros(n_test_modes) + + if _NUMBA_AVAILABLE: + ref_real = ref_shape.real + ref_imag = ref_shape.imag + + for test_idx in range(n_test_modes): + test_shape = test_shapes[:, test_idx] + mac_row[test_idx] = _fast_mac_computation( + ref_real, ref_imag, test_shape.real, test_shape.imag + ) + else: + for test_idx in range(n_test_modes): + test_shape = test_shapes[:, test_idx] + + numerator = (ref_shape.conj().T @ test_shape) * (test_shape.conj().T @ ref_shape) + denominator = (ref_shape.conj().T @ ref_shape) * (test_shape.conj().T @ test_shape) + + if np.abs(denominator) == 0: + mac_row[test_idx] = 0.0 + else: + mac_row[test_idx] = np.real(numerator / denominator) + + return mac_row + + if _JOBLIB_AVAILABLE and n_jobs != 1: + mac_rows = Parallel(n_jobs=n_jobs)( + delayed(compute_mac_row)(ref_idx) for ref_idx in range(n_ref_modes) + ) + else: + mac_rows = [compute_mac_row(ref_idx) for ref_idx in range(n_ref_modes)] + + return np.array(mac_rows) + + +def get_performance_info() -> Dict[str, bool]: + """Get information about available performance optimizations. + + Returns: + Dictionary with availability of optimization libraries + """ + return { + "numba_available": _NUMBA_AVAILABLE, + "joblib_available": _JOBLIB_AVAILABLE, + "parallel_processing": _JOBLIB_AVAILABLE, + "jit_compilation": _NUMBA_AVAILABLE + } + + +# ============================================================================= +# MATLAB Assessment Functions Implementation (arassess.m and corrassess.m) +# ============================================================================= + +def arassess( + time_vector: Optional[npt.NDArray[np.floating]], + signal_data: npt.NDArray[np.floating], + ar_order: int, + block_params: Tuple[int, int], + median_filter_length: Optional[int] = None, + channel_names: Optional[List[str]] = None +) -> AssessmentResult: + """ + Auto-detect problematic channels using AR modeling predictability and participation. + + Direct Python port of MATLAB arassess.m function for analyzing channel quality through: + - Predictability (DAF): How well each channel can be predicted from AR model + - Participation (PRT): How much each channel contributes to AR model + - RMS levels: Signal power levels + + Args: + time_vector: Time vector, shape (N,) or None for sample indices + signal_data: Signal data matrix, shape (N, M) where M is channels + ar_order: AR model order (scalar, past lag horizon p) + block_params: Block parameters (block_size, block_stride) + median_filter_length: Optional median filter length (>=3 to apply) + channel_names: Optional channel names list + + Returns: + AssessmentResult containing channel health assessment + """ + N, M = signal_data.shape + block_size, block_stride = block_params + + # Normalize channels to unity RMS (following MATLAB arassess.m exactly) + Y = signal_data.copy().astype(np.float64) + + # Compute RY matrix and RMS values exactly as in MATLAB + RY = (Y.T @ Y) / N + rmsy = np.sqrt(np.diag(RY)) + + # Normalize each channel by its RMS + for cc in range(M): + if rmsy[cc] > 0: + Y[:, cc] = Y[:, cc] / rmsy[cc] + + # Check rank of covariance matrix + if np.linalg.matrix_rank(RY) != M: + print("Warning: signal data is not full rank") + + # Apply median filtering if requested (MATLAB logic) + if median_filter_length is not None and median_filter_length >= 3: + from scipy.signal import medfilt + for cc in range(M): + Y[:, cc] = medfilt(Y[:, cc], kernel_size=median_filter_length) + + # Set up channel names matching MATLAB format + if channel_names is None: + channel_names = [f"#{cc+1}" for cc in range(M)] + else: + if len(channel_names) != M: + raise ValueError("Channel names length must match number of channels") + # Add channel numbers like MATLAB version + channel_names = [f"{name} (#{cc+1})" for cc, name in enumerate(channel_names)] + + # Block analysis setup (MATLAB variable names) + t1vec = np.arange(0, N - block_size + 1, block_stride) + NBlock = len(t1vec) + + if NBlock == 0: + raise ValueError("No complete blocks can be formed with given parameters") + + # Time vector for blocks + if time_vector is not None and len(time_vector) > 0: + t = np.zeros(NBlock) + for jj in range(NBlock): + t1, t2 = t1vec[jj], t1vec[jj] + block_size - 1 + t[jj] = (time_vector[t1] + time_vector[t2]) / 2 + else: + t = np.arange(1, NBlock + 1) # MATLAB 1-based indexing style + + # Initialize results matrices (matching MATLAB variable names) + DAF = np.zeros((M, NBlock)) # Deviation Accounted For (predictability) + PRT = np.zeros((M, NBlock)) # Participation/relevance + RMS = np.zeros((M, NBlock)) # RMS levels + + for jj in range(NBlock): + t1 = t1vec[jj] + t2 = t1 + block_size + + # Process block - fit AR model using subvar equivalent + try: + block_data = Y[t1:t2, :] + H, Ry, Re = _subvar_matlab(block_data, ar_order) + + # Deviation accounted for (predictability) - exact MATLAB formula + dafjj = np.maximum(0, 1 - np.sqrt(np.diag(Re) / np.diag(Ry))) + + # Participation for each channel - exact MATLAB logic + prtjj = np.zeros(M) + for ii in range(M): + # Extract H coefficients for channel ii across all lags: H(:,ii:M:(p*M)) + # H is shape (M, M * ar_order), channel ii appears at positions ii, ii+M, ii+2*M, etc. + channel_cols = np.arange(ii, M * ar_order, M) + prtjj[ii] = np.linalg.norm(H[:, channel_cols], 'fro') + + # Normalize participation to sum to M (exact MATLAB formula: M*prtjj/sum(prtjj)) + prt_sum = np.sum(prtjj) + if prt_sum > 0: + prtjj = M * prtjj / prt_sum + + # Store results + DAF[:, jj] = dafjj + PRT[:, jj] = prtjj + RMS[:, jj] = np.sqrt(np.diag(Ry)) * rmsy + + except (np.linalg.LinAlgError, ValueError) as e: + print(f"Warning: AR model failed for block {jj}: {e}") + # Set default values for failed blocks + DAF[:, jj] = 0.0 + PRT[:, jj] = 1.0 + RMS[:, jj] = rmsy + + # Compute median values (meddpr in MATLAB) + meddpr = np.column_stack([ + np.median(DAF, axis=1), # Median predictability + np.median(PRT, axis=1), # Median participation + np.median(RMS, axis=1) # Median RMS + ]) + + return AssessmentResult( + predictability=DAF, + participation=PRT, + rms_values=RMS, + median_scores=meddpr, + time_vector=t, + channel_names=channel_names, + block_parameters=block_params, + ar_lag=ar_order + ) + + +def corrassess( + time_vector: Optional[npt.NDArray[np.floating]], + signal_data: npt.NDArray[np.floating], + block_params: Tuple[int, int], + channel_names: Optional[List[str]] = None +) -> CorrelationAssessmentResult: + """ + Simple correlation map output to assess channel correlation patterns. + + Direct Python port of MATLAB corrassess.m for analyzing correlations between channels + to identify problematic or redundant channels. + + Args: + time_vector: Time vector, shape (N,) or None for sample indices + signal_data: Signal data matrix, shape (N, M) where M is number of channels + block_params: Block parameters (block_size, block_stride) + channel_names: Optional channel names list + + Returns: + CorrelationAssessmentResult containing correlation analysis results + """ + N, M = signal_data.shape + block_size, block_stride = block_params + + # Check data rank (following MATLAB corrassess.m) + RY = (signal_data.T @ signal_data) / N + if np.linalg.matrix_rank(RY) != M: + print("Warning: signal data is not full rank") + + # Set up channel names matching MATLAB format + if channel_names is None: + channel_names = [f"#{cc+1}" for cc in range(M)] + else: + if len(channel_names) != M: + raise ValueError("Channel names length must match number of channels") + # Add channel numbers like MATLAB version + channel_names = [f"{name} (#{cc+1})" for cc, name in enumerate(channel_names)] + + # Block analysis setup (MATLAB variable names) + t1vec = np.arange(0, N - block_size + 1, block_stride) + NBlock = len(t1vec) + + if NBlock == 0: + raise ValueError("No complete blocks can be formed with given parameters") + + # Time vector for blocks + if time_vector is not None and len(time_vector) > 0: + t = np.zeros(NBlock) + for jj in range(NBlock): + t1, t2 = t1vec[jj], t1vec[jj] + block_size - 1 + t[jj] = (time_vector[t1] + time_vector[t2]) / 2 + else: + t = np.arange(1, NBlock + 1) # MATLAB 1-based indexing style + + # Correlation map + CMAP = np.zeros((M, NBlock)) + + for jj in range(NBlock): + t1 = t1vec[jj] + t2 = t1 + block_size + + # Process block + Z = signal_data[t1:t2, :] + C = np.abs(_corrloc_matlab(Z)) + + # Average correlation with other channels: c=(sum(C,2)-diag(C))/(M-1) + c = (np.sum(C, axis=1) - np.diag(C)) / (M - 1) + CMAP[:, jj] = c + + # Median correlation across blocks: C=median(CMAP,2) + median_correlation = np.median(CMAP, axis=1) + + return CorrelationAssessmentResult( + median_correlation=median_correlation, + correlation_map=CMAP, + time_blocks=t, + channel_names=channel_names + ) + + +def _subvar_matlab(y: npt.NDArray[np.floating], p: int) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """ + Fit multivariate AR model using least squares. + + Direct Python translation of the subvar function from arassess.m + + Args: + y: Signal data, shape (N, ny) + p: AR model order (number of lags) + + Returns: + H: AR coefficient matrix, shape (ny, ny*p) + Ry: Signal covariance matrix, shape (ny, ny) + Re: Residual covariance matrix, shape (ny, ny) + """ + N, ny = y.shape + Nk = N - p # Number of effective data points + nyp = ny * p # Total number of AR parameters per equation + + # Create matrices following MATLAB indexing + Z = y.T # Transpose for easier column access + Zp = np.zeros((nyp, Nk)) # Past data matrix + Y = np.zeros((ny, Nk)) # Current data matrix + + # Build lagged data matrix exactly as in MATLAB + for kk in range(p, N): # MATLAB: for kk=(p+1):N + # Extract past p samples: (kk-1):-1:(kk-p) in MATLAB 1-indexing + # zp=reshape(Z(:,(kk-1):-1:(kk-p)),nyp,1); + past_cols = [] + for lag_idx in range(p): + past_cols.append(Z[:, kk - 1 - lag_idx]) + + # Reshape and store (MATLAB: reshape(..., nyp, 1)) + zp = np.concatenate(past_cols) + Zp[:, kk - p] = zp + Y[:, kk - p] = y[kk, :].T + + # Solve AR equations: Y = H * Zp + E + # MATLAB: Rzp=Zp*Zp'; H=(Y*Zp')/Rzp; + Rzp = Zp @ Zp.T + H = (Y @ Zp.T) @ np.linalg.pinv(Rzp) # AR coefficient matrix + + # Compute covariances (exact MATLAB formulas) + Ry = (Y @ Y.T) / Nk # Signal covariance: Ry=(Y*Y')/Nk; + E = Y - H @ Zp # Residuals: E=Y-H*Zp; + Re = (E @ E.T) / Nk # Residual covariance: Re=(E*E')/Nk; + + return H, Ry, Re + + +def _corrloc_matlab(Z: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: + """ + Compute local correlation matrix for a data block. + + Direct Python translation of corrloc function from corrassess.m + + Args: + Z: Data block, shape (n, m) + + Returns: + r: Correlation matrix, shape (m, m) + """ + n, m = Z.shape + + # Normalize each column (MATLAB loop logic) + Z_norm = np.zeros_like(Z) + for jj in range(m): + # MATLAB: Z(:,jj)=Z(:,jj)-mean(Z(:,jj)); Z(:,jj)=Z(:,jj)/norm(Z(:,jj)); + col = Z[:, jj] - np.mean(Z[:, jj]) + norm_col = np.linalg.norm(col) + if norm_col > 0: + Z_norm[:, jj] = col / norm_col + else: + Z_norm[:, jj] = col + + # MATLAB: r=(Z'*Z); + r = Z_norm.T @ Z_norm + return r + + +# ============================================================================= +# Additional MATLAB Signal Processing Functions +# ============================================================================= + +def yfilt( + signal_data: npt.NDArray[np.floating], + Ts: Optional[float] = None, + filter_type: Literal["LP", "HP", "BP", "BS"] = "LP", + frequency_band: Union[float, Tuple[float, float]] = 0.1, + filter_order: int = 2, + use_filtfilt: bool = True +) -> npt.NDArray[np.floating]: + """ + Filter columns of signal data independently as time-series. + + Python port of MATLAB yfilt.m that applies Butterworth filters to + multichannel signals with zero-phase distortion option. + + Args: + signal_data: Signal data where each column is a channel, shape (N, M) + Ts: Sample period (None for normalized frequencies) + filter_type: Filter type ("LP", "HP", "BP", "BS") + frequency_band: Cutoff frequency (scalar) or band (tuple) + filter_order: Filter order (default 2) + use_filtfilt: Use zero-phase filtering (default True) + + Returns: + Filtered signal data, same shape as input + """ + from scipy.signal import butter, filtfilt, lfilter + + N, M = signal_data.shape + + # Handle frequency normalization + if Ts is not None: + Fs = 1.0 / Ts + if isinstance(frequency_band, tuple): + fband = np.array(frequency_band) * (2.0 / Fs) + else: + fband = frequency_band * (2.0 / Fs) + else: + fband = frequency_band + + # Check Nyquist constraint + if isinstance(fband, np.ndarray): + if np.any(fband >= 1.0): + raise ValueError("Frequencies specified beyond Nyquist limit") + else: + if fband >= 1.0: + raise ValueError("Frequency specified beyond Nyquist limit") + + # Design filter + if filter_type.upper() == "LP": + b, a = butter(filter_order, fband, btype='low') + elif filter_type.upper() == "HP": + b, a = butter(filter_order, fband, btype='high') + elif filter_type.upper() == "BP": + b, a = butter(filter_order, fband, btype='band') + elif filter_type.upper() == "BS": + b, a = butter(filter_order, fband, btype='bandstop') + else: + raise ValueError(f"Unknown filter type: {filter_type}") + + # Apply filter to each column + filtered_data = np.zeros_like(signal_data) + + if use_filtfilt: + for jj in range(M): + filtered_data[:, jj] = filtfilt(b, a, signal_data[:, jj]) + else: + for jj in range(M): + filtered_data[:, jj] = lfilter(b, a, signal_data[:, jj]) + + return filtered_data + + +def yinterpolate( + time_vector: npt.NDArray[np.floating], + signal_data: npt.NDArray[np.floating], + new_sample_rate: float +) -> Tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]: + """ + Upsample signal to new sample rate with anti-aliasing filtering. + + Python port of MATLAB yinterpolate.m that upsamples signals using + interpolation followed by anti-aliasing low-pass filtering. + + Args: + time_vector: Original time vector, shape (N,) + signal_data: Signal data, shape (N, M) + new_sample_rate: New sample rate (must be higher than original) + + Returns: + new_time: New time vector + new_signal: Interpolated and filtered signal data + """ + from scipy.interpolate import PchipInterpolator + + # Compute original sample rate + Ts = np.mean(np.diff(time_vector)) + Fs = 1.0 / Ts + + if new_sample_rate <= Fs: + raise ValueError("Only upsampling is supported") + + # Create new time vector + Ts_new = 1.0 / new_sample_rate + t_min, t_max = time_vector.min(), time_vector.max() + new_time = np.arange(t_min, t_max + Ts_new, Ts_new) + + # Interpolate each channel + N_new, M = len(new_time), signal_data.shape[1] + new_signal = np.zeros((N_new, M)) + + # Upsample factor and anti-aliasing cutoff + r = new_sample_rate / Fs + w_cut = 1.0 / r + + for cc in range(M): + # Interpolate using PCHIP (similar to MATLAB's pchip) + interpolator = PchipInterpolator(time_vector, signal_data[:, cc]) + new_signal[:, cc] = interpolator(new_time) + + # Apply anti-aliasing filter + new_signal[:, cc] = yfilt( + new_signal[:, cc:cc+1], + Ts=Ts_new, + filter_type="LP", + frequency_band=w_cut, + filter_order=8 + ).ravel() + + return new_time, new_signal + + +def qplot_data( + data: npt.NDArray[np.floating], + quantiles: Optional[npt.NDArray[np.floating]] = None, + column_labels: Optional[List[str]] = None +) -> Tuple[npt.NDArray[np.floating], List[str]]: + """ + Compute quantile statistics for plotting. + + Python port of MATLAB qplot.m that computes quantile statistics + for each column of data matrix. Returns data for plotting rather + than creating plots directly. + + Args: + data: Data matrix, shape (n, m) + quantiles: Quantile values to compute (default: [0.01, 0.1, 0.33, 0.5, 0.67, 0.9, 0.99]) + column_labels: Column labels (generated if None) + + Returns: + quantile_values: Quantile values for each column, shape (7, m) + labels: Column labels + """ + if quantiles is None: + quantiles = np.array([0.01, 0.1, 1/3, 0.5, 2/3, 0.9, 0.99]) + + if len(quantiles) != 7: + raise ValueError("quantiles must have 7 elements") + + if np.any(quantiles >= 1) or np.any(quantiles <= 0): + raise ValueError("all quantile values must be in (0,1)") + + quantiles = np.sort(quantiles) + + n, m = data.shape + + # Generate labels if not provided + if column_labels is None: + column_labels = [f"#{i+1}" for i in range(m)] + elif len(column_labels) != m: + raise ValueError("Number of labels must match number of columns") + + # Compute quantiles for each column + quantile_values = np.zeros((7, m)) + + for ii in range(m): + column_data = np.sort(data[:, ii]) + indices = np.round(quantiles * (n - 1)).astype(int) + quantile_values[:, ii] = column_data[indices] + + return quantile_values, column_labels + + +# ============================================================================= +# FFT Spectral Analysis Functions (Missing from MATLAB) +# ============================================================================= + +def fftspec( + time_vector: Optional[NDArray[np.number]], + signal_data: NDArray[np.number], + block_params: Tuple[int, int], + nfft: Union[int, Tuple[int, int]], + reduced_dim: int = 0, + contour_levels: int = 25 +) -> FFTSpectralResult: + """ + Multi-channel FFT-based power spectral density analysis with optional random projection. + + This function computes block-based FFT spectral analysis of multivariate time series, + with optional dimensionality reduction via random projection for computational efficiency. + + Args: + time_vector: Time vector or None for sample indices + signal_data: Input data matrix, shape (n_samples, n_channels) + block_params: (block_size, block_stride) for windowing + nfft: FFT length, or (nfft, smoothing_span) + reduced_dim: If > 0, project to this dimension before FFT + contour_levels: Number of contour levels for visualization + + Returns: + FFTSpectralResult containing power spectral density and frequency information + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*t)[:, np.newaxis] + >>> result = fftspec(t, y, (256, 128), 512) + >>> print(f"PSD shape: {result.P.shape}") + """ + if signal_data.ndim != 2: + raise ValueError(f"signal_data must be 2D, got shape {signal_data.shape}") + + n_samples, n_channels = signal_data.shape + block_size, block_stride = block_params + + # Handle nfft parameter + if isinstance(nfft, (tuple, list)): + nfft_len, n_smooth = nfft[:2] + else: + nfft_len, n_smooth = nfft, 1 + + # Handle time vector + if time_vector is None: + time_vector = np.arange(n_samples) + fs = 1.0 + else: + time_vector = np.asarray(time_vector).flatten() + if len(time_vector) != n_samples: + raise ValueError("Length of time_vector does not match signal_data") + fs = 1.0 / (time_vector[1] - time_vector[0]) if len(time_vector) > 1 else 1.0 + + # Apply random projection if requested + if reduced_dim > 0 and reduced_dim < n_channels: + projection_matrix = np.random.randn(reduced_dim, n_channels) + signal_data = (projection_matrix @ signal_data.T).T + n_channels = reduced_dim + + # Compute block-based FFT + block_starts = np.arange(0, n_samples - block_size + 1, block_stride) + n_blocks = len(block_starts) + + # Initialize output arrays + freq_bins = fftfreq(nfft_len, 1/fs)[:nfft_len//2] # Positive frequencies only + n_freq = len(freq_bins) + psd_matrix = np.zeros((n_freq, n_blocks)) + block_times = np.zeros(n_blocks) + + # Process each block + for i, start_idx in enumerate(block_starts): + end_idx = start_idx + block_size + block_data = signal_data[start_idx:end_idx, :] + block_times[i] = time_vector[start_idx + block_size // 2] + + # Remove mean from each channel + block_data = block_data - np.mean(block_data, axis=0) + + # Compute FFT for each channel and aggregate + block_psd = np.zeros(n_freq) + for ch in range(n_channels): + # Zero-pad if necessary + if block_size < nfft_len: + padded_data = np.zeros(nfft_len) + padded_data[:block_size] = block_data[:, ch] + else: + padded_data = block_data[:nfft_len, ch] + + # Compute FFT and PSD + fft_data = fft(padded_data) + channel_psd = np.abs(fft_data[:n_freq])**2 / (fs * nfft_len) + block_psd += channel_psd + + # Average across channels + psd_matrix[:, i] = block_psd / n_channels + + # Apply smoothing if requested + if n_smooth > 1: + from scipy.ndimage import uniform_filter1d + psd_matrix = uniform_filter1d(psd_matrix, size=n_smooth, axis=0) + + return FFTSpectralResult( + P=psd_matrix, + F=freq_bins, + T=block_times, + nfft=nfft_len, + block_size=block_size, + n_blocks=n_blocks, + fs=fs + ) + + +def fftspec1( + signal: NDArray[np.number], + nfft: int, + overlap: float = 0.5, + window: str = 'hann' +) -> Tuple[NDArray[np.floating], NDArray[np.floating]]: + """ + Single-channel FFT spectral analysis using Welch's method. + + Args: + signal: Single-channel time series + nfft: FFT length + overlap: Overlap fraction (0 to 1) + window: Window function name + + Returns: + Tuple of (frequencies, power spectral density) + """ + from scipy.signal import welch + + if signal.ndim != 1: + signal = signal.flatten() + + nperseg = nfft + noverlap = int(nperseg * overlap) + + frequencies, psd = welch( + signal, + fs=1.0, + window=window, + nperseg=nperseg, + noverlap=noverlap, + nfft=nfft + ) + + return frequencies, psd + + +def fftspecwin( + signal_data: NDArray[np.number], + window_params: Dict[str, Union[str, int, float]] +) -> FFTSpectralResult: + """ + Windowed FFT spectral analysis with various window functions. + + Args: + signal_data: Input data matrix, shape (n_samples, n_channels) + window_params: Dictionary with window parameters: + - 'type': Window type ('hann', 'hamming', 'blackman', etc.) + - 'nfft': FFT length + - 'overlap': Overlap fraction + - 'detrend': Detrending method ('linear', 'constant', None) + + Returns: + FFTSpectralResult containing windowed spectral analysis + """ + from scipy.signal import spectrogram + + window_type = window_params.get('type', 'hann') + nfft = window_params.get('nfft', 1024) + overlap = window_params.get('overlap', 0.5) + detrend_method = window_params.get('detrend', 'constant') + + if signal_data.ndim == 1: + signal_data = signal_data[:, np.newaxis] + + n_samples, n_channels = signal_data.shape + noverlap = int(nfft * overlap) + + # Compute spectrogram for each channel and aggregate + total_psd = None + frequencies = None + times = None + + for ch in range(n_channels): + f, t, psd = spectrogram( + signal_data[:, ch], + fs=1.0, + window=window_type, + nperseg=nfft, + noverlap=noverlap, + nfft=nfft, + detrend=detrend_method + ) + + if total_psd is None: + total_psd = psd + frequencies = f + times = t + else: + total_psd += psd + + # Average across channels + total_psd /= n_channels + + return FFTSpectralResult( + P=total_psd, + F=frequencies, + T=times, + nfft=nfft, + block_size=nfft, + n_blocks=len(times), + fs=1.0 + ) + + +def zmfftspec( + signal_data: NDArray[np.number], + remove_mean: bool = True, + remove_trend: bool = False +) -> Tuple[NDArray[np.floating], NDArray[np.floating]]: + """ + Zero-mean FFT spectral analysis. + + Args: + signal_data: Input signal, shape (n_samples,) or (n_samples, n_channels) + remove_mean: Whether to remove DC component + remove_trend: Whether to remove linear trend + + Returns: + Tuple of (frequencies, power spectral density) + """ + if signal_data.ndim == 1: + signal_data = signal_data[:, np.newaxis] + + n_samples, n_channels = signal_data.shape + + # Preprocess signal + processed_data = signal_data.copy() + + if remove_trend: + from scipy.signal import detrend + for ch in range(n_channels): + processed_data[:, ch] = detrend(processed_data[:, ch]) + elif remove_mean: + processed_data = processed_data - np.mean(processed_data, axis=0) + + # Compute FFT + fft_data = fft(processed_data, axis=0) + frequencies = fftfreq(n_samples)[:n_samples//2] + + # Compute PSD and average across channels + psd = np.mean(np.abs(fft_data[:n_samples//2, :])**2, axis=1) / n_samples + + return frequencies, psd + + +# ============================================================================= +# Signal Conditioning Functions (Missing from MATLAB) +# ============================================================================= + +def yintegrate( + signal: NDArray[np.number], + method: str = 'trapz', + initial_value: float = 0.0 +) -> NDArray[np.floating]: + """ + Numerical integration of time series using various methods. + + Args: + signal: Input signal to integrate + method: Integration method ('trapz', 'simpson', 'cumsum') + initial_value: Initial condition for integration + + Returns: + Integrated signal + """ + if method == 'trapz': + from scipy.integrate import cumulative_trapezoid + return initial_value + cumulative_trapezoid(signal, initial=0) + elif method == 'simpson': + from scipy.integrate import simpson + # For cumulative Simpson's rule, we need to compute incrementally + integrated = np.zeros_like(signal) + integrated[0] = initial_value + for i in range(1, len(signal)): + if i == 1: + integrated[i] = integrated[i-1] + (signal[i] + signal[i-1]) / 2 + else: + integrated[i] = integrated[i-1] + simpson(signal[i-1:i+1]) + return integrated + elif method == 'cumsum': + return initial_value + np.cumsum(signal) + else: + raise ValueError(f"Unknown integration method: {method}") + + +def ydecimate( + signal: NDArray[np.number], + decimation_factor: int, + filter_order: int = 8, + filter_type: str = 'iir' +) -> NDArray[np.floating]: + """ + Decimation with anti-aliasing filtering. + + Args: + signal: Input signal to decimate + decimation_factor: Factor by which to reduce sampling rate + filter_order: Order of anti-aliasing filter + filter_type: Type of filter ('iir' or 'fir') + + Returns: + Decimated signal + """ + from scipy.signal import decimate + + return decimate(signal, decimation_factor, n=filter_order, ftype=filter_type) + + +def yresample( + signal: NDArray[np.number], + original_rate: float, + target_rate: float, + method: str = 'linear' +) -> NDArray[np.floating]: + """ + Resample signal to new sampling rate with interpolation. + + Args: + signal: Input signal to resample + original_rate: Original sampling rate + target_rate: Target sampling rate + method: Interpolation method ('linear', 'cubic', 'nearest') + + Returns: + Resampled signal + """ + from scipy.interpolate import interp1d + + n_original = len(signal) + n_target = int(n_original * target_rate / original_rate) + + # Create time vectors + t_original = np.linspace(0, n_original / original_rate, n_original) + t_target = np.linspace(0, n_original / original_rate, n_target) + + # Interpolate + interpolator = interp1d(t_original, signal, kind=method, + bounds_error=False, fill_value='extrapolate') + + return interpolator(t_target) + + +def yaddgauss( + signal: NDArray[np.number], + noise_level: float, + random_seed: Optional[int] = None +) -> NDArray[np.floating]: + """ + Add Gaussian noise to signal. + + Args: + signal: Input signal + noise_level: Standard deviation of noise relative to signal + random_seed: Random seed for reproducibility + + Returns: + Signal with added Gaussian noise + """ + if random_seed is not None: + np.random.seed(random_seed) + + signal_std = np.std(signal) + noise = np.random.normal(0, noise_level * signal_std, signal.shape) + + return signal + noise + + +# ============================================================================= +# Advanced Spectral Analysis Functions (Missing from MATLAB) +# ============================================================================= + +def fdmspec1( + time_vector: Optional[NDArray[np.number]], + signal_data: NDArray[np.number], + block_params: Tuple[int, int], + fdm_params: Tuple[int, int, int, float, int], + threshold: Union[float, Tuple[float, float]] = 0.99, + random_seed: Optional[int] = None +) -> Dict: + """ + Block-based finite difference modal (FDM) frequency analysis. + + This function implements FDM-like frequency analysis using compressed sampling + and random projections, following the MATLAB fdmspec1.m algorithm. + + Args: + time_vector: Time vector or None for sample indices + signal_data: Input data matrix, shape (n_samples, n_channels) + block_params: (block_size, block_stride) for windowing + fdm_params: (r, d1, d2, alpha, K) where: + - r: Reduced dimension (<0 for orthonormal, >0 for random, 0 for none) + - d1, d2: FDM eigenproblem sizes + - alpha: Regularization parameter (<0 for auto, 0 for none) + - K: Maximum number of frequencies to return + threshold: Threshold value(s) close to 1 (e.g., 0.99) + random_seed: Random seed for reproducibility + + Returns: + Dictionary containing FDM analysis results in eigspec-compatible format + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*5*t)[:, np.newaxis] + >>> result = fdmspec1(t, y, (256, 128), (-10, 50, 75, -1, 10)) + """ + if signal_data.ndim != 2: + raise ValueError(f"signal_data must be 2D, got shape {signal_data.shape}") + if len(block_params) != 2: + raise ValueError("block_params must have 2 elements") + if len(fdm_params) != 5: + raise ValueError("fdm_params must have 5 elements") + + if isinstance(threshold, (int, float)): + threshold = (threshold, threshold) + elif len(threshold) != 2: + raise ValueError("threshold must be scalar or 2-element tuple") + + if not all(0 <= t <= 1 for t in threshold): + raise ValueError("threshold values must be between 0 and 1") + + if random_seed is not None: + np.random.seed(random_seed) + + n_samples, n_channels = signal_data.shape + block_size, block_stride = block_params + r, d1, d2, alpha, K = fdm_params + + # Handle time vector + if time_vector is None: + time_vector = np.arange(n_samples) + time_step = -1 + else: + time_vector = np.asarray(time_vector).flatten() + if len(time_vector) != n_samples: + raise ValueError("Length of time_vector does not match signal_data") + time_step = time_vector[1] - time_vector[0] if len(time_vector) > 1 else 1.0 + + # Calculate block parameters + block_starts = np.arange(0, n_samples - block_size + 1, block_stride) + n_blocks = len(block_starts) + + # Initialize results structure (mock-up to match rndspec format) + block_results = [] + + for i, start_idx in enumerate(block_starts): + end_idx = start_idx + block_size + block_data = signal_data[start_idx:end_idx, :].copy() + + # Demean block + block_data = block_data - np.mean(block_data, axis=0) + + # Apply random projection if specified + if r != 0: + if r < 0: + # Orthonormal projection + projection_matrix = np.linalg.qr(np.random.randn(n_channels, -r))[0].T + projected_data = (projection_matrix @ block_data.T).T + else: + # Random projection + projection_matrix = np.random.randn(r, n_channels) + projected_data = (projection_matrix @ block_data.T).T + else: + projection_matrix = None + projected_data = block_data + + # FDM frequency analysis using eigenvalue-based method + frequencies, amplitudes = _fdm_eigenanalysis( + projected_data, d1, d2, alpha, K, threshold[0] + ) + + # Create mock modal report structure + block_result = { + 'modal_report': { + 'frequencies': frequencies, + 'amplitudes': amplitudes, + 'lambda_vals': np.exp(1j * 2 * np.pi * frequencies) if len(frequencies) > 0 else np.array([]), + 'imode': list(range(len(frequencies))) if len(frequencies) > 0 else [] + }, + 'shape_estimates': None, + 'projection_matrix': projection_matrix, + 'processing_time': 0.0, # Would be filled in real implementation + 'demean_block': True, + 'time_step': time_step, + 'time_slice': (start_idx, end_idx), + 'centre_time': float((time_vector[start_idx] + time_vector[end_idx-1]) / 2), + 'filter_time': float(time_vector[end_idx-1]) + } + + block_results.append(block_result) + + return { + 'block_results': block_results, + 'fdm_params': fdm_params, + 'block_params': block_params, + 'threshold': threshold, + 'routine': 'fdmspec1' + } + + +def _fdm_eigenanalysis( + data: NDArray[np.number], + d1: int, + d2: int, + alpha: float, + max_frequencies: int, + threshold: float +) -> Tuple[NDArray[np.floating], NDArray[np.floating]]: + """ + Core FDM eigenanalysis for frequency extraction. + + This is a simplified implementation of the FDM eigenanalysis procedure. + """ + n_samples, n_channels = data.shape + + if n_samples < max(d1, d2) + 1: + return np.array([]), np.array([]) + + # Build Hankel-like matrices for FDM + # This is a simplified version - full implementation would be more complex + + # Use autocorrelation-based approach as approximation + frequencies = [] + amplitudes = [] + + # For each channel, find dominant frequencies using autocorrelation + for ch in range(n_channels): + signal = data[:, ch] + + # Compute autocorrelation + autocorr = np.correlate(signal, signal, mode='full') + autocorr = autocorr[len(autocorr)//2:] + + # Find peaks in autocorrelation (simplified frequency detection) + if len(autocorr) > 10: + from scipy.signal import find_peaks + peaks, properties = find_peaks(autocorr[1:], height=threshold * np.max(autocorr)) + + # Convert peak positions to frequencies + for peak in peaks[:max_frequencies//n_channels]: + if peak > 0: + freq = 1.0 / (peak + 1) # Simplified frequency estimation + amp = autocorr[peak + 1] + frequencies.append(freq) + amplitudes.append(amp) + + # Sort by amplitude and keep top frequencies + if len(frequencies) > 0: + freq_amp_pairs = sorted(zip(frequencies, amplitudes), key=lambda x: x[1], reverse=True) + frequencies = np.array([f for f, a in freq_amp_pairs[:max_frequencies]]) + amplitudes = np.array([a for f, a in freq_amp_pairs[:max_frequencies]]) + else: + frequencies = np.array([]) + amplitudes = np.array([]) + + return frequencies, amplitudes + + +def kdftspec( + time_vector: Optional[NDArray[np.number]], + signal_data: NDArray[np.number], + block_params: Tuple[int, int], + dft_params: Union[int, Tuple[int, int]], + reduced_dim: int = 0, + k_folds: int = 8 +) -> FFTSpectralResult: + """ + K-fold DFT spectral analysis with robustified pruning of bursty data. + + This function implements k-fold cross-validated DFT analysis that acts as + a de-ELMing filter for spectrograms, following MATLAB kdftspec.m. + + Args: + time_vector: Time vector or None for sample indices + signal_data: Input data matrix, shape (n_samples, n_channels) + block_params: (block_size, block_stride) for windowing + dft_params: DFT length, or (dft_length, smoothing_span) + reduced_dim: Reduced dimension for random projection (0 for none) + k_folds: Number of folds for robustification + + Returns: + FFTSpectralResult with robustified spectral analysis + + Example: + >>> t = np.linspace(0, 10, 1000) + >>> y = np.sin(2*np.pi*t)[:, np.newaxis] + >>> result = kdftspec(t, y, (512, 128), (2048, 3), 10, 8) + """ + if signal_data.ndim != 2: + raise ValueError(f"signal_data must be 2D, got shape {signal_data.shape}") + if len(block_params) != 2: + raise ValueError("block_params must have 2 elements") + + n_samples, n_channels = signal_data.shape + block_size, block_stride = block_params + + # Handle DFT parameters + if isinstance(dft_params, (int, float)): + dft_length, smoothing_span = int(dft_params), 1 + else: + dft_length, smoothing_span = dft_params[:2] + + # Handle time vector + if time_vector is None: + time_vector = np.arange(n_samples) + fs = 1.0 + else: + time_vector = np.asarray(time_vector).flatten() + if len(time_vector) != n_samples: + raise ValueError("Length of time_vector does not match signal_data") + fs = 1.0 / (time_vector[1] - time_vector[0]) if len(time_vector) > 1 else 1.0 + + # Apply random projection if specified + if reduced_dim > 0 and reduced_dim < n_channels: + projection_matrix = np.random.randn(reduced_dim, n_channels) + signal_data = (projection_matrix @ signal_data.T).T + n_channels = reduced_dim + + # Compute block-based DFT with k-fold robustification + block_starts = np.arange(0, n_samples - block_size + 1, block_stride) + n_blocks = len(block_starts) + + # Initialize output arrays + freq_bins = fftfreq(dft_length, 1/fs)[:dft_length//2] + n_freq = len(freq_bins) + psd_matrix = np.zeros((n_freq, n_blocks)) + block_times = np.zeros(n_blocks) + + # Process each block with k-fold robustification + for i, start_idx in enumerate(block_starts): + end_idx = start_idx + block_size + block_data = signal_data[start_idx:end_idx, :] + block_times[i] = time_vector[start_idx + block_size // 2] + + # Remove mean from each channel + block_data = block_data - np.mean(block_data, axis=0) + + # K-fold cross-validation for robustification + fold_psds = [] + fold_size = block_size // k_folds + + for fold in range(k_folds): + fold_start = fold * fold_size + fold_end = min(fold_start + fold_size, block_size) + + if fold_end <= fold_start: + continue + + fold_data = block_data[fold_start:fold_end, :] + + # Compute FFT for this fold + fold_psd = np.zeros(n_freq) + for ch in range(n_channels): + # Zero-pad if necessary + if len(fold_data) < dft_length: + padded_data = np.zeros(dft_length) + padded_data[:len(fold_data)] = fold_data[:, ch] + else: + padded_data = fold_data[:dft_length, ch] + + # Compute FFT and PSD + fft_data = fft(padded_data) + channel_psd = np.abs(fft_data[:n_freq])**2 / (fs * dft_length) + fold_psd += channel_psd + + fold_psds.append(fold_psd / n_channels) + + # Robustified PSD using median or mean of folds + if len(fold_psds) > 0: + fold_psds = np.array(fold_psds) + # Use median for robustification (removes outliers/bursts) + psd_matrix[:, i] = np.median(fold_psds, axis=0) + else: + psd_matrix[:, i] = 0 + + # Apply smoothing if requested + if smoothing_span > 1: + if smoothing_span > 0: + # Moving average smoothing + from scipy.ndimage import uniform_filter1d + psd_matrix = uniform_filter1d(psd_matrix, size=smoothing_span, axis=0) + else: + # Median filter smoothing + from scipy.ndimage import median_filter + psd_matrix = median_filter(psd_matrix, size=(-smoothing_span, 1)) + + return FFTSpectralResult( + P=psd_matrix, + F=freq_bins, + T=block_times, + nfft=dft_length, + block_size=block_size, + n_blocks=n_blocks, + fs=fs + ) \ No newline at end of file diff --git a/src/tokeye/eigspec/utils/subspace_identification.py b/src/tokeye/eigspec/utils/subspace_identification.py new file mode 100644 index 0000000..dee9bbf --- /dev/null +++ b/src/tokeye/eigspec/utils/subspace_identification.py @@ -0,0 +1,586 @@ +""" +Subspace identification algorithms for eigspec package. + +This module provides stochastic subspace identification (SSI) algorithms for system identification: +- Covariance-driven SSI (SSI-COV) with block Hankel matrices +- Canonical correlation analysis SSI (SSI-CCA) for robust identification +- State-space model extraction and validation + +Based on the MATLAB eigspec toolbox subspace identification functions: +- ssi1ca.m - Covariance-driven stochastic subspace identification +- ssicca.m - Canonical correlation analysis subspace identification +- ssi1cax.m - Extended SSI with cross-validation +- kfoldcov.m - K-fold cross-validation for covariance estimation +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import numpy.typing as npt +from numpy.typing import NDArray + + +@dataclass +class SubspaceIdentificationResult: + """ + Container for subspace identification analysis results. + + Attributes: + state_matrix: State transition matrix A, shape (n_states, n_states) + output_matrix: Output matrix C, shape (n_outputs, n_states) + singular_values: Singular values from SVD decomposition + models: List of models for multiple orders (if applicable) + """ + state_matrix: Optional[NDArray[np.floating]] + output_matrix: Optional[NDArray[np.floating]] + singular_values: NDArray[np.floating] + models: Optional[List['StateSpaceModel']] + + +@dataclass +class StateSpaceModel: + """ + Container for individual state-space model. + + Attributes: + state_matrix: State transition matrix A + output_matrix: Output matrix C + kalman_gain: Kalman gain matrix K (for CCA variant) + """ + state_matrix: NDArray[np.floating] + output_matrix: NDArray[np.floating] + kalman_gain: Optional[NDArray[np.floating]] = None + + +def covariance_driven_ssi( + data: NDArray[np.number], + identification_params: List[int] +) -> SubspaceIdentificationResult: + """ + Covariance-driven stochastic subspace identification using classical SVD approach. + + This function implements the classical covariance-driven SSI algorithm that extracts + state-space models (A, C) from multichannel time series data based on the extended + observability matrix. + + Args: + data: Multichannel data where each column is a channel time-series, + shape (n_samples, n_channels) + identification_params: Analysis parameters [future, past, order1, order2, ...] + - future: Number of future samples in block-Hankel matrix + - past: Number of past samples in block-Hankel matrix + - order1, order2, ...: System orders to estimate + + Returns: + SubspaceIdentificationResult containing: + - state_matrix, output_matrix for single order case + - List of models for multiple orders + - Singular values from SVD + + Raises: + ValueError: If system order exceeds theoretical limit or parameters are invalid + + Example: + >>> # Generate synthetic 2-channel oscillatory data + >>> t = np.linspace(0, 10, 1000) + >>> data = np.column_stack([np.sin(2*np.pi*t), np.cos(2*np.pi*t)]) + >>> params = [10, 10, 2, 4] # future=10, past=10, orders=[2,4] + >>> result = covariance_driven_ssi(data, params) + >>> print(f"Found {len(result.models)} models") + """ + if not isinstance(data, np.ndarray): + raise TypeError("data must be a numpy array") + if data.ndim != 2: + raise ValueError(f"data must be 2D, got shape {data.shape}") + if not isinstance(identification_params, list) or len(identification_params) < 3: + raise ValueError("identification_params must be a list with at least 3 elements [future, past, order1, ...]") + + n_samples, n_channels = data.shape + future_horizon = identification_params[0] + past_horizon = identification_params[1] + system_orders = identification_params[2:] + + if len(system_orders) == 0: + raise ValueError("System order list is empty") + if any(order > n_channels * future_horizon for order in system_orders): + max_order = max(system_orders) + raise ValueError(f"Order request {max_order} exceeds theoretical limit n_channels*future = {n_channels*future_horizon}") + if future_horizon <= 0 or past_horizon <= 0: + raise ValueError("Future and past horizons must be positive integers") + if n_samples <= past_horizon + future_horizon: + raise ValueError(f"Data length {n_samples} too short for past+future = {past_horizon+future_horizon}") + + # Construct block-Hankel data matrix + n_data_columns = n_samples - past_horizon - future_horizon + 1 + n_data_rows = n_channels * (past_horizon + future_horizon) + hankel_matrix = np.zeros((n_data_rows, n_data_columns)) + + for col_idx in range(n_data_columns): + # Extract data block from time col_idx to col_idx+past+future-1. + # Rows must be time-block-major ([all channels at t0, all channels + # at t1, ...]) to match the n_channels-strided slicing below (and + # the MATLAB original); flattening the (time, channel) block in C + # order gives exactly that. + data_block = data[col_idx:col_idx+past_horizon+future_horizon, :] + hankel_matrix[:, col_idx] = data_block.flatten() + + # Split into past and future components + past_data = hankel_matrix[:n_channels*past_horizon, :] + future_data = hankel_matrix[n_channels*past_horizon:, :] + + # Calculate the past-to-future projection matrix + cross_covariance = (future_data @ past_data.T) / n_data_columns + + # SVD decomposition + # Visualize SVD Output? + left_singular_vectors, singular_values, right_singular_vectors_T = np.linalg.svd(cross_covariance, full_matrices=False) + + if len(system_orders) == 1: + # Single model case + order = system_orders[0] + + # Construct extended observability matrix + observability_matrix = left_singular_vectors[:, :order] @ np.diag(np.sqrt(singular_values[:order])) + + # Extract system matrices + output_matrix = observability_matrix[:n_channels, :] + # MATLAB: A = O(1:(m*(f-1)),:) \ O((m+1):(m*f),:). + # lstsq(O1, O2)[0] is exactly pinv(O1) @ O2 — the same as backslash; + # no transpose is needed (or correct) here. + state_matrix = np.linalg.lstsq( + observability_matrix[:n_channels*(future_horizon-1), :], + observability_matrix[n_channels:n_channels*future_horizon, :], + rcond=None + )[0] + + return SubspaceIdentificationResult( + state_matrix=state_matrix, + output_matrix=output_matrix, + singular_values=singular_values, + models=None + ) + + else: + # Multiple model case + models = [] + for order in system_orders: + # Construct extended observability matrix for this order + observability_matrix = left_singular_vectors[:, :order] @ np.diag(np.sqrt(singular_values[:order])) + + # Extract system matrices + output_matrix = observability_matrix[:n_channels, :] + # Same as the single-model case: lstsq already matches backslash. + state_matrix = np.linalg.lstsq( + observability_matrix[:n_channels*(future_horizon-1), :], + observability_matrix[n_channels:n_channels*future_horizon, :], + rcond=None + )[0] + + models.append(StateSpaceModel(state_matrix=state_matrix, output_matrix=output_matrix)) + + return SubspaceIdentificationResult( + state_matrix=None, + output_matrix=None, + singular_values=singular_values, + models=models + ) + + +def canonical_correlation_ssi( + data: NDArray[np.number], + identification_params: List[int], + compute_residual: bool = False +) -> SubspaceIdentificationResult: + """ + Stochastic subspace identification using canonical correlation analysis. + + This function implements the CCA-based SSI algorithm which can provide better + numerical conditioning compared to basic covariance-driven SSI. + + Args: + data: Multichannel data, shape (n_samples, n_channels) + identification_params: Analysis parameters [future, past, order1, order2, ...] + compute_residual: If True, compute residual error covariance + + Returns: + SubspaceIdentificationResult containing models with A, K, C matrices and singular values + + Example: + >>> data = np.random.randn(1000, 3) # 3-channel data + >>> params = [15, 15, 3] # future=15, past=15, order=3 + >>> result = canonical_correlation_ssi(data, params) + """ + if not isinstance(data, np.ndarray): + raise TypeError("data must be a numpy array") + if data.ndim != 2: + raise ValueError(f"data must be 2D, got shape {data.shape}") + + n_samples, n_channels = data.shape + future_horizon = identification_params[0] + past_horizon = identification_params[1] + system_orders = identification_params[2:] + + if len(system_orders) == 0: + raise ValueError("System order list is empty") + if any(order > n_channels * future_horizon for order in system_orders): + max_order = max(system_orders) + raise ValueError(f"Order request exceeds theoretical limit n_channels*future = {n_channels*future_horizon}") + + # Construct block-Hankel data matrix + n_data_columns = n_samples - past_horizon - future_horizon + 1 + n_data_rows = n_channels * (past_horizon + future_horizon) + hankel_matrix = np.zeros((n_data_rows, n_data_columns)) + + for col_idx in range(n_data_columns): + data_block = data[col_idx:col_idx+past_horizon+future_horizon, :].T + hankel_matrix[:, col_idx] = data_block.flatten() + + # Split into past and future components + past_data = hankel_matrix[:n_channels*past_horizon, :] + future_data = hankel_matrix[n_channels*past_horizon:, :] + + # Calculate covariance matrices for CCA + future_covariance = (future_data @ future_data.T) / n_data_columns + past_covariance = (past_data @ past_data.T) / n_data_columns + cross_covariance = (future_data @ past_data.T) / n_data_columns + + # Compute CCA weights using matrix square roots + U1, S1, Vh1 = np.linalg.svd(future_covariance) + inverse_sqrt_future_cov = Vh1.T @ np.diag(1.0 / np.sqrt(S1)) @ Vh1 + + U1, S1, Vh1 = np.linalg.svd(past_covariance) + inverse_sqrt_past_cov = Vh1.T @ np.diag(1.0 / np.sqrt(S1)) @ Vh1 + + # CCA-weighted matrix + cca_matrix = inverse_sqrt_future_cov @ cross_covariance @ inverse_sqrt_past_cov + left_vectors, singular_values, right_vectors_T = np.linalg.svd(cca_matrix, full_matrices=False) + + if len(system_orders) == 1: + # Single model case + order = system_orders[0] + + # Compute state sequence + state_sequence = (right_vectors_T[:order, :] @ inverse_sqrt_past_cov) @ past_data + output_sequence = data[past_horizon:n_samples-future_horizon+1, :].T + + # Estimate output matrix C + output_matrix = output_sequence @ np.linalg.pinv(state_sequence) + + # Estimate A and K matrices + n_time_steps = state_sequence.shape[1] + augmented_regression_matrix = np.vstack([ + state_sequence[:, :n_time_steps-1], + data[past_horizon:n_samples-future_horizon, :].T + ]) + next_states = state_sequence[:, 1:n_time_steps] + + system_kalman_matrix = next_states @ np.linalg.pinv(augmented_regression_matrix) + state_matrix = system_kalman_matrix[:, :order] + kalman_gain = system_kalman_matrix[:, order:order+n_channels] + final_state_matrix = state_matrix + kalman_gain @ output_matrix + + model = StateSpaceModel( + state_matrix=final_state_matrix, + output_matrix=output_matrix, + kalman_gain=kalman_gain + ) + return SubspaceIdentificationResult( + state_matrix=final_state_matrix, + output_matrix=output_matrix, + singular_values=singular_values, + models=[model] + ) + + else: + # Multiple model case + models = [] + for order in system_orders: + # Compute state sequence + state_sequence = (right_vectors_T[:order, :] @ inverse_sqrt_past_cov) @ past_data + output_sequence = data[past_horizon:n_samples-future_horizon+1, :].T + + # Estimate matrices + output_matrix = output_sequence @ np.linalg.pinv(state_sequence) + + n_time_steps = state_sequence.shape[1] + augmented_regression_matrix = np.vstack([ + state_sequence[:, :n_time_steps-1], + data[past_horizon:n_samples-future_horizon, :].T + ]) + next_states = state_sequence[:, 1:n_time_steps] + + system_kalman_matrix = next_states @ np.linalg.pinv(augmented_regression_matrix) + state_matrix = system_kalman_matrix[:, :order] + kalman_gain = system_kalman_matrix[:, order:order+n_channels] + final_state_matrix = state_matrix + kalman_gain @ output_matrix + + models.append(StateSpaceModel( + state_matrix=final_state_matrix, + output_matrix=output_matrix, + kalman_gain=kalman_gain + )) + + return SubspaceIdentificationResult( + state_matrix=None, + output_matrix=None, + singular_values=singular_values, + models=models + ) + + +def ssi1ca( + signal_data: npt.NDArray[np.floating], + params: Tuple[int, int, Union[int, List[int]]] +) -> Dict[str, Any]: + """ + Covariance-driven stochastic subspace identification. + + Python port of MATLAB ssi1ca.m that performs basic covariance-driven SSI + using extended observability matrix with no particular SVD weighting. + + Args: + signal_data: Multichannel data where columns are channels, shape (N, m) + params: [future, past, order(s)] parameters + + Returns: + Dictionary containing A, C matrices and singular values + """ + N, m = signal_data.shape + future, past = params[0], params[1] + + # Handle order specification + if len(params) == 3: + if isinstance(params[2], (list, np.ndarray)): + orders = params[2] + else: + orders = [params[2]] + else: + orders = list(params[2:]) + + if len(orders) == 0: + raise ValueError("System order list is empty") + + if np.any(np.array(orders) > m * future): + raise ValueError("Order request exceeds m*f") + + # Build data matrix + n_cols = N - past - future + 1 + n_rows = m * (past + future) + D = np.zeros((n_rows, n_cols)) + + for kk in range(n_cols): + block = signal_data[kk:kk+past+future, :].T.ravel() + D[:, kk] = block + + # Split into past and future + Yp = D[:m*past, :] + Yf = D[m*past:, :] + + # Compute past-to-future projection + Rfp = (Yf @ Yp.T) / n_cols + U, S, Vt = np.linalg.svd(Rfp, full_matrices=False) + sigma = np.diag(S) + + if len(orders) == 1: + # Single model with given order + n = orders[0] + O = U[:, :n] @ np.diag(np.sqrt(sigma[:n])) + + # Split observability matrix to get C and A + C = O[:m, :] + if O.shape[0] > m: + O1 = O[:-m, :] + O2 = O[m:, :] + # Solve for A: O2 = A * O1 + A = np.linalg.lstsq(O1.T, O2.T, rcond=None)[0].T + else: + A = np.zeros((n, n)) + + return { + 'A': A, + 'C': C, + 'sigm': sigma, + 'order': n + } + else: + # Multiple orders - return observability matrices for each + results = {} + for i, n in enumerate(orders): + O = U[:, :n] @ np.diag(np.sqrt(sigma[:n])) + C = O[:m, :] + + if O.shape[0] > m: + O1 = O[:-m, :] + O2 = O[m:, :] + A = np.linalg.lstsq(O1.T, O2.T, rcond=None)[0].T + else: + A = np.zeros((n, n)) + + results[f'order_{n}'] = { + 'A': A, + 'C': C, + 'sigm': sigma, + 'order': n + } + + results['sigm'] = sigma + return results + + +def ssicca( + signal_data: npt.NDArray[np.floating], + params: Tuple[int, int, Union[int, List[int]]], + compute_kalman: bool = True +) -> Dict[str, Any]: + """ + Stochastic subspace identification using canonical correlation analysis. + + Python port of MATLAB ssicca.m that performs SSI using CCA weighting + for improved numerical properties and noise handling. + + Args: + signal_data: Multichannel data where columns are channels, shape (N, m) + params: [future, past, order(s)] parameters + compute_kalman: Whether to compute Kalman gain matrix + + Returns: + Dictionary containing A, C, K matrices and singular values + """ + N, m = signal_data.shape + future, past = params[0], params[1] + + # Handle order specification + if len(params) == 3: + if isinstance(params[2], (list, np.ndarray)): + orders = params[2] + else: + orders = [params[2]] + else: + orders = list(params[2:]) + + if len(orders) == 0: + raise ValueError("System order list is empty") + + if np.any(np.array(orders) > m * future): + raise ValueError("Order request exceeds m*f") + + # Build data matrix + n_cols = N - past - future + 1 + n_rows = m * (past + future) + D = np.zeros((n_rows, n_cols)) + + for kk in range(n_cols): + block = signal_data[kk:kk+past+future, :].T.ravel() + D[:, kk] = block + + # Split into past and future + Yp = D[:m*past, :] + Yf = D[m*past:, :] + + # Compute covariance matrices + Rff = (Yf @ Yf.T) / n_cols + Rpp = (Yp @ Yp.T) / n_cols + Rfp = (Yf @ Yp.T) / n_cols + + # CCA weighting: compute inverse square roots + try: + # Eigendecomposition for matrix square root inverse + U1, S1, Vt1 = np.linalg.svd(Rff) + inv_sqrt_Rff = Vt1.T @ np.diag(1.0 / np.sqrt(S1)) @ Vt1 + + U2, S2, Vt2 = np.linalg.svd(Rpp) + inv_sqrt_Rpp = Vt2.T @ np.diag(1.0 / np.sqrt(S2)) @ Vt2 + + # CCA matrix + M = inv_sqrt_Rff @ Rfp @ inv_sqrt_Rpp + + except np.linalg.LinAlgError: + # Fallback to regularized version if matrices are singular + reg_eps = 1e-10 + Rff_reg = Rff + reg_eps * np.eye(Rff.shape[0]) + Rpp_reg = Rpp + reg_eps * np.eye(Rpp.shape[0]) + + U1, S1, Vt1 = np.linalg.svd(Rff_reg) + inv_sqrt_Rff = Vt1.T @ np.diag(1.0 / np.sqrt(S1)) @ Vt1 + + U2, S2, Vt2 = np.linalg.svd(Rpp_reg) + inv_sqrt_Rpp = Vt2.T @ np.diag(1.0 / np.sqrt(S2)) @ Vt2 + + M = inv_sqrt_Rff @ Rfp @ inv_sqrt_Rpp + + # SVD of CCA matrix + U, S, Vt = np.linalg.svd(M, full_matrices=False) + sigma = np.diag(S) + + if len(orders) == 1: + # Single model + n = orders[0] + + # Construct extended observability matrix with CCA weighting + O = inv_sqrt_Rff @ U[:, :n] @ np.diag(np.sqrt(sigma[:n])) + + # Extract C and A matrices + C = O[:m, :] + if O.shape[0] > m: + O1 = O[:-m, :] + O2 = O[m:, :] + A = np.linalg.lstsq(O1.T, O2.T, rcond=None)[0].T + else: + A = np.zeros((n, n)) + + result = { + 'A': A, + 'C': C, + 'sigm': sigma, + 'order': n + } + + # Compute Kalman gain if requested + if compute_kalman and O.shape[0] > m: + try: + residual = O2 - A @ O1 + Ree = (residual @ residual.T) / residual.shape[1] + K = np.linalg.lstsq(C.T, np.eye(m), rcond=None)[0].T + result['K'] = K + result['Ree'] = Ree + except np.linalg.LinAlgError: + pass + + return result + + else: + # Multiple orders + results = {} + for i, n in enumerate(orders): + O = inv_sqrt_Rff @ U[:, :n] @ np.diag(np.sqrt(sigma[:n])) + C = O[:m, :] + + if O.shape[0] > m: + O1 = O[:-m, :] + O2 = O[m:, :] + A = np.linalg.lstsq(O1.T, O2.T, rcond=None)[0].T + else: + A = np.zeros((n, n)) + + order_result = { + 'A': A, + 'C': C, + 'sigm': sigma, + 'order': n + } + + if compute_kalman and O.shape[0] > m: + try: + residual = O2 - A @ O1 + Ree = (residual @ residual.T) / residual.shape[1] + K = np.linalg.lstsq(C.T, np.eye(m), rcond=None)[0].T + order_result['K'] = K + order_result['Ree'] = Ree + except np.linalg.LinAlgError: + pass + + results[f'order_{n}'] = order_result + + results['sigm'] = sigma + return results diff --git a/src/tokeye/eigspec/utils/utils.py b/src/tokeye/eigspec/utils/utils.py new file mode 100644 index 0000000..a399d00 --- /dev/null +++ b/src/tokeye/eigspec/utils/utils.py @@ -0,0 +1,157 @@ +""" +General utility functions for eigspec package. + +This module provides common utility functions used throughout the eigspec package: +- Array and matrix manipulation utilities +- Mathematical helper functions +- Data validation and processing utilities +- Common computational routines + +Based on utility functions scattered throughout the MATLAB eigspec toolbox: +- Various helper functions in eigspec_mmain.m +- Mathematical utilities in subspace identification functions +- Array processing utilities in modal analysis functions +- General computational helpers used across the MATLAB codebase +""" + +from typing import Tuple +import numpy as np +from numpy.typing import ArrayLike, NDArray + +def demean(x: NDArray[np.number], axis: int = 0, inplace: bool = False) -> NDArray[np.number]: + """ + Remove mean from input data along specified axis. + + Args: + x: Input data array + axis: Axis along which to compute mean (default: 0) + inplace: If True, modify array in place (default: False) + + Returns: + Demeaned data array of same shape as input + + Example: + >>> x = np.array([[1, 2], [3, 4]]) + >>> demean(x) + array([[-1, -1], + [ 1, 1]]) + """ + if not isinstance(x, np.ndarray): + raise TypeError("Input must be a numpy array") + + if inplace: + x -= np.mean(x, axis=axis, keepdims=True) + return x + return x - np.mean(x, axis=axis, keepdims=True) + +def complex_to_real(z: NDArray[np.complexfloating]) -> Tuple[NDArray[np.floating], NDArray[np.floating]]: + """ + Convert complex array to real and imaginary parts. + + Args: + z: Complex input array + + Returns: + Tuple of (real_part, imag_part) + + Example: + >>> z = np.array([1+2j, 3+4j]) + >>> real, imag = complex_to_real(z) + >>> real + array([1., 3.]) + >>> imag + array([2., 4.]) + """ + if not np.iscomplexobj(z): + raise TypeError("Input must be a complex array") + return np.real(z), np.imag(z) + +def real_to_complex(real_part: ArrayLike, imag_part: ArrayLike) -> NDArray[np.complexfloating]: + """ + Convert real and imaginary parts to complex array. + + Args: + real_part: Real part array + imag_part: Imaginary part array + + Returns: + Complex array + + Example: + >>> real = np.array([1, 3]) + >>> imag = np.array([2, 4]) + >>> real_to_complex(real, imag) + array([1.+2.j, 3.+4.j]) + """ + real_arr = np.asarray(real_part) + imag_arr = np.asarray(imag_part) + if real_arr.shape != imag_arr.shape: + raise ValueError(f"Shape mismatch: real {real_arr.shape} != imag {imag_arr.shape}") + return real_arr + 1j * imag_arr + +def validate_dimensions(*arrays: NDArray) -> None: + """ + Validate that input arrays have compatible dimensions. + + Args: + *arrays: Variable number of numpy arrays to validate + + Raises: + ValueError: If dimensions are incompatible + TypeError: If any input is not a numpy array + + Example: + >>> a = np.zeros((2, 3)) + >>> b = np.ones((2, 3)) + >>> validate_dimensions(a, b) # No error + >>> c = np.zeros((3, 2)) + >>> validate_dimensions(a, c) # Raises ValueError + """ + if not arrays: + return + + if not all(isinstance(arr, np.ndarray) for arr in arrays): + raise TypeError("All inputs must be numpy arrays") + + shape = arrays[0].shape + mismatched = [(i, arr.shape) for i, arr in enumerate(arrays[1:], 1) + if arr.shape != shape] + + if mismatched: + details = [f"array[{i}]: {s}" for i, s in mismatched] + raise ValueError(f"Array shapes must match. Found shape {shape} but got:\n" + + "\n".join(details)) + +def is_positive_definite(matrix: NDArray[np.number], rtol: float = 1e-5) -> bool: + """ + Check if a matrix is positive definite using Cholesky decomposition. + + Args: + matrix: Square matrix to check + rtol: Relative tolerance for numerical stability checks + + Returns: + True if matrix is positive definite + + Example: + >>> A = np.array([[2, -1], [-1, 2]]) # Positive definite + >>> is_positive_definite(A) + True + >>> B = np.array([[1, 2], [2, 1]]) # Not positive definite + >>> is_positive_definite(B) + False + """ + if not isinstance(matrix, np.ndarray): + raise TypeError("Input must be a numpy array") + + if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]: + raise ValueError(f"Expected square matrix, got shape {matrix.shape}") + + if not np.allclose(matrix, matrix.T, rtol=rtol): + return False + + try: + np.linalg.cholesky(matrix) + return True + except np.linalg.LinAlgError: + return False \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/__init__.py b/src/tokeye/eigspec/vis/__init__.py new file mode 100644 index 0000000..b60a6d7 --- /dev/null +++ b/src/tokeye/eigspec/vis/__init__.py @@ -0,0 +1,135 @@ +""" +Visualization module for eigspec package. + +This module provides plotting and visualization functions for spectral analysis results: +- Spectral plots for frequency-domain analysis +- Time series plotting with modal overlays +- Clustering visualization and pattern recognition plots +- Modal shape and pattern visualization +- Utility plots for analysis diagnostics + +Based on the MATLAB eigspec toolbox visualization functions: +- view_pcaspec_cluster.m - Cluster visualization and analysis +- view_pcaspec_results.m - Main results plotting +- view_pcaspec_prototypes.m - Prototype-based visualization +- view_pcaspec_cluster_pattern.m - Pattern estimation and plotting +- view_pcaspec_cluster_nearness.m - Cluster similarity visualization +- Various plotting utilities throughout the MATLAB toolbox +""" + +# Use conditional imports to handle missing modules gracefully +try: + from .spectral_plots import ( + SpectralPlotOptions, + plot_eigenvalue_evolution, + plot_frequency_time, + plot_phase_time, + plot_rms_time, + plot_spectral_summary, + view_pcaspec, + view_pcaspec_results, + ) + _SPECTRAL_AVAILABLE = True +except ImportError: + _SPECTRAL_AVAILABLE = False + +try: + from .modal_plots import ( + ModalPlotOptions, + plot_mode_shapes, + plot_mode_shape_2d, + plot_mode_shape_polar, + plot_shape_vectors, + plot_array_geometry, + ) + _MODAL_AVAILABLE = True +except ImportError: + _MODAL_AVAILABLE = False + +try: + from .clustering_plots import ( + ClusterPlotOptions, + plot_clustering_results, + plot_cluster_similarity_matrix, + plot_cluster_medoids, + plot_mac_similarity, + plot_soft_clustering, + ) + _CLUSTERING_AVAILABLE = True +except ImportError: + _CLUSTERING_AVAILABLE = False + +try: + from .time_series_plots import ( + TimeSeriesPlotOptions, + plot_time_traces, + plot_prototype_traces, + plot_filtered_traces, + plot_multi_channel_overlay, + ) + _TIME_SERIES_AVAILABLE = True +except ImportError: + _TIME_SERIES_AVAILABLE = False + +try: + from .utility_plots import ( + plot_quantiles, + plot_statistical_summary, + create_colormap, + setup_figure_style, + ) + _UTILITY_AVAILABLE = True +except ImportError: + _UTILITY_AVAILABLE = False + +# Build __all__ list dynamically based on available imports +__all__ = [] + +if _SPECTRAL_AVAILABLE: + __all__.extend([ + 'SpectralPlotOptions', + 'plot_eigenvalue_evolution', + 'plot_frequency_time', + 'plot_phase_time', + 'plot_rms_time', + 'plot_spectral_summary', + 'view_pcaspec', + 'view_pcaspec_results', + ]) + +if _MODAL_AVAILABLE: + __all__.extend([ + 'ModalPlotOptions', + 'plot_mode_shapes', + 'plot_mode_shape_2d', + 'plot_mode_shape_polar', + 'plot_shape_vectors', + 'plot_array_geometry', + ]) + +if _CLUSTERING_AVAILABLE: + __all__.extend([ + 'ClusterPlotOptions', + 'plot_clustering_results', + 'plot_cluster_similarity_matrix', + 'plot_cluster_medoids', + 'plot_mac_similarity', + 'plot_soft_clustering', + ]) + +if _TIME_SERIES_AVAILABLE: + __all__.extend([ + 'TimeSeriesPlotOptions', + 'plot_time_traces', + 'plot_prototype_traces', + 'plot_filtered_traces', + 'plot_multi_channel_overlay', + ]) + +if _UTILITY_AVAILABLE: + __all__.extend([ + 'plot_quantiles', + 'plot_statistical_summary', + 'create_colormap', + 'setup_figure_style', + ]) \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/clustering_plots.py b/src/tokeye/eigspec/vis/clustering_plots.py new file mode 100644 index 0000000..8b97959 --- /dev/null +++ b/src/tokeye/eigspec/vis/clustering_plots.py @@ -0,0 +1,462 @@ +""" +Clustering visualization functions for eigspec package. + +This module provides specialized plotting functions for clustering analysis results: +- Cluster scatter plots with frequency-time coloring +- MAC-based similarity and distance visualization +- Medoid and prototype shape plotting +- Cluster validation and quality metrics visualization + +Based on the MATLAB eigspec toolbox clustering visualization functions: +- view_pcaspec_cluster.m - Main cluster visualization and analysis +- view_pcaspec_unsupervised.m - Unsupervised clustering plots +- view_pcaspec_cluster_nearness.m - Cluster similarity visualization +- view_pcaspec_cluster_pattern.m - Pattern estimation and plotting +- redraw_clustering() - Cluster result plotting utilities +""" + +from typing import Optional, List, Tuple, Union, Dict, Any +import numpy as np +import numpy.typing as npt +import matplotlib.pyplot as plt +from matplotlib import colormaps +import matplotlib.colors as mcolors +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from dataclasses import dataclass + + +@dataclass +class ClusterPlotOptions: + """Configuration options for clustering analysis plots. + + Attributes: + figsize: Figure size (width, height) in inches + dpi: Figure DPI for resolution + fontsize: Base font size for text elements + title_fontsize: Font size for plot titles + label_fontsize: Font size for axis labels + legend_fontsize: Font size for legend text + grid: Whether to show grid lines + colormap: Colormap name for cluster colors + marker_size: Size of scatter plot markers + line_width: Width of plot lines + alpha: Transparency level (0-1) + save_format: Default format for saving figures + """ + figsize: Tuple[float, float] = (12, 8) + dpi: int = 100 + fontsize: int = 12 + title_fontsize: int = 14 + label_fontsize: int = 12 + legend_fontsize: int = 10 + grid: bool = True + colormap: str = 'tab10' + marker_size: float = 6.0 + line_width: float = 1.5 + alpha: float = 0.8 + save_format: str = 'png' + + +def plot_clustering_results( + time: npt.NDArray[np.floating], + frequencies: npt.NDArray[np.floating], + cluster_labels: npt.NDArray[np.integer], + medoid_indices: Optional[npt.NDArray[np.integer]] = None, + amplitudes: Optional[npt.NDArray[np.floating]] = None, + options: Optional[ClusterPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot clustering results in time-frequency space. + + Args: + time: Time vector + frequencies: Frequency values + cluster_labels: Cluster assignment for each point + medoid_indices: Indices of cluster medoids + amplitudes: Optional amplitude values for marker sizing + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ClusterPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Get unique clusters + unique_clusters = np.unique(cluster_labels) + n_clusters = len(unique_clusters) + + # Create colors for clusters + if n_clusters <= 10: + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_clusters)) + else: + colors = colormaps.get_cmap('hsv')(np.linspace(0, 1, n_clusters)) + + # Plot each cluster + for i, cluster_id in enumerate(unique_clusters): + mask = cluster_labels == cluster_id + + # Determine marker size + if amplitudes is not None: + sizes = options.marker_size**2 * (1 + amplitudes[mask]) + else: + sizes = options.marker_size**2 + + ax.scatter(time[mask], frequencies[mask], + c=[colors[i]], s=sizes, alpha=options.alpha, + label=f'Cluster {cluster_id}') + + # Highlight medoids if provided + if medoid_indices is not None: + ax.scatter(time[medoid_indices], frequencies[medoid_indices], + c='black', marker='x', s=(options.marker_size*2)**2, + alpha=1.0, linewidths=3, label='Medoids') + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Frequency', fontsize=options.label_fontsize) + ax.set_title('Clustering Results', fontsize=options.title_fontsize) + + # Add legend with reasonable number of entries + if n_clusters <= 20: + ax.legend(fontsize=options.legend_fontsize, bbox_to_anchor=(1.05, 1), + loc='upper left') + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_cluster_similarity_matrix( + similarity_matrix: npt.NDArray[np.floating], + cluster_labels: Optional[npt.NDArray[np.integer]] = None, + method: str = 'similarity', + options: Optional[ClusterPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot similarity or distance matrix as heatmap. + + Args: + similarity_matrix: Similarity or distance matrix + cluster_labels: Optional cluster labels for ordering + method: Type of matrix ('similarity' or 'distance') + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ClusterPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Order by clusters if labels provided + if cluster_labels is not None: + order = np.argsort(cluster_labels) + ordered_matrix = similarity_matrix[order][:, order] + ordered_labels = cluster_labels[order] + else: + ordered_matrix = similarity_matrix + ordered_labels = None + + # Choose colormap based on method + if method == 'similarity': + cmap = 'Blues' + label = 'Similarity' + else: + cmap = 'Reds' + label = 'Distance' + + # Create heatmap + im = ax.imshow(ordered_matrix, cmap=cmap, alpha=options.alpha, + aspect='auto', origin='lower') + + # Add colorbar + cbar = plt.colorbar(im, ax=ax) + cbar.set_label(label, fontsize=options.label_fontsize) + + # Add cluster boundaries if labels provided + if ordered_labels is not None: + # Find cluster boundaries + boundaries = [] + current_cluster = ordered_labels[0] + for i, label in enumerate(ordered_labels[1:], 1): + if label != current_cluster: + boundaries.append(i - 0.5) + current_cluster = label + + # Draw boundary lines + for boundary in boundaries: + ax.axhline(boundary, color='red', linewidth=2, alpha=0.7) + ax.axvline(boundary, color='red', linewidth=2, alpha=0.7) + + ax.set_xlabel('Data Point Index', fontsize=options.label_fontsize) + ax.set_ylabel('Data Point Index', fontsize=options.label_fontsize) + ax.set_title(f'{label} Matrix', fontsize=options.title_fontsize) + + return fig, ax + + +def plot_cluster_medoids( + shape_vectors: npt.NDArray[np.complex128], + cluster_labels: npt.NDArray[np.integer], + medoid_indices: npt.NDArray[np.integer], + coordinates: Optional[npt.NDArray[np.floating]] = None, + options: Optional[ClusterPlotOptions] = None +) -> Figure: + """Plot cluster medoids and their shape vectors. + + Args: + shape_vectors: Complex shape vectors (n_channels x n_vectors) + cluster_labels: Cluster assignment for each vector + medoid_indices: Indices of cluster medoids + coordinates: Optional sensor coordinates for spatial plotting + options: Plot configuration options + + Returns: + Figure object with subplots + """ + if options is None: + options = ClusterPlotOptions() + + unique_clusters = np.unique(cluster_labels) + n_clusters = len(unique_clusters) + n_channels = shape_vectors.shape[0] + + # Determine subplot layout + n_cols = min(3, n_clusters) + n_rows = (n_clusters + n_cols - 1) // n_cols + + fig, axes = plt.subplots(n_rows, n_cols, figsize=(options.figsize[0]*n_cols/3, + options.figsize[1]*n_rows/2), + dpi=options.dpi) + + # Handle single subplot case + if n_clusters == 1: + axes = [axes] + elif n_rows == 1: + axes = axes.flatten() + else: + axes = axes.flatten() + + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_clusters)) + + for i, cluster_id in enumerate(unique_clusters): + if i >= len(axes): + break + + medoid_idx = medoid_indices[i] + medoid_vector = shape_vectors[:, medoid_idx] + + # Normalize medoid + medoid_vector = medoid_vector / np.sqrt(np.sum(np.abs(medoid_vector)**2)) + + if coordinates is not None: + # Spatial plot + axes[i].scatter(coordinates[:, 0], np.real(medoid_vector), + c='blue', marker='o', s=options.marker_size**2, + alpha=options.alpha, label='Real') + axes[i].scatter(coordinates[:, 1], np.imag(medoid_vector), + c='red', marker='s', s=options.marker_size**2, + alpha=options.alpha, label='Imaginary') + axes[i].set_xlabel('Coordinate') + else: + # Channel index plot + channels = np.arange(n_channels) + axes[i].plot(channels, np.real(medoid_vector), 'b-o', + linewidth=options.line_width, markersize=options.marker_size, + alpha=options.alpha, label='Real') + axes[i].plot(channels, np.imag(medoid_vector), 'r-s', + linewidth=options.line_width, markersize=options.marker_size, + alpha=options.alpha, label='Imaginary') + axes[i].set_xlabel('Channel Index') + + axes[i].set_ylabel('Amplitude') + axes[i].set_title(f'Cluster {cluster_id} Medoid', fontsize=options.title_fontsize) + axes[i].legend(fontsize=options.legend_fontsize-2) + + if options.grid: + axes[i].grid(True, alpha=0.3) + + # Hide unused subplots + for i in range(n_clusters, len(axes)): + axes[i].set_visible(False) + + plt.tight_layout() + return fig + + +def plot_mac_similarity( + mac_matrix: npt.NDArray[np.floating], + threshold: Optional[float] = None, + cluster_labels: Optional[npt.NDArray[np.integer]] = None, + options: Optional[ClusterPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot MAC (Modal Assurance Criterion) similarity matrix. + + Args: + mac_matrix: MAC similarity matrix + threshold: Optional threshold for highlighting + cluster_labels: Optional cluster labels for ordering + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ClusterPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Order by clusters if labels provided + if cluster_labels is not None: + order = np.argsort(cluster_labels) + ordered_matrix = mac_matrix[order][:, order] + else: + ordered_matrix = mac_matrix + + # Create heatmap + im = ax.imshow(ordered_matrix, cmap='viridis', alpha=options.alpha, + aspect='auto', origin='lower', vmin=0, vmax=1) + + # Add threshold contour if provided + if threshold is not None: + contour = ax.contour(ordered_matrix, levels=[threshold], colors='red', + linewidths=2, alpha=0.8) + ax.clabel(contour, inline=True, fontsize=options.fontsize-2) + + # Add colorbar + cbar = plt.colorbar(im, ax=ax) + cbar.set_label('MAC Value', fontsize=options.label_fontsize) + + ax.set_xlabel('Mode Index', fontsize=options.label_fontsize) + ax.set_ylabel('Mode Index', fontsize=options.label_fontsize) + title = 'MAC Similarity Matrix' + if threshold is not None: + title += f' (threshold = {threshold:.2f})' + ax.set_title(title, fontsize=options.title_fontsize) + + return fig, ax + + +def plot_soft_clustering( + time: npt.NDArray[np.floating], + frequencies: npt.NDArray[np.floating], + membership_matrix: npt.NDArray[np.floating], + method: str = 'interpolated', + options: Optional[ClusterPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot soft clustering results with membership probabilities. + + Args: + time: Time vector + frequencies: Frequency values + membership_matrix: Membership probabilities (n_points x n_clusters) + method: Soft clustering visualization method + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ClusterPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + n_clusters = membership_matrix.shape[1] + + if method == 'interpolated': + # Use membership probabilities as colors + for i in range(n_clusters): + membership = membership_matrix[:, i] + scatter = ax.scatter(time, frequencies, c=membership, + cmap='viridis', s=options.marker_size**2, + alpha=options.alpha, vmin=0, vmax=1) + + cbar = plt.colorbar(scatter, ax=ax) + cbar.set_label('Membership Probability', fontsize=options.label_fontsize) + + elif method == 'nearest_neighbor': + # Use k-nearest neighbor style visualization + from scipy.spatial.distance import cdist + + # Create grid for interpolation + time_grid = np.linspace(time.min(), time.max(), 50) + freq_grid = np.linspace(frequencies.min(), frequencies.max(), 50) + TIME, FREQ = np.meshgrid(time_grid, freq_grid) + + # Find nearest neighbors and interpolate membership + points = np.column_stack([time, frequencies]) + grid_points = np.column_stack([TIME.ravel(), FREQ.ravel()]) + + distances = cdist(grid_points, points) + nearest_indices = np.argmin(distances, axis=1) + + for i in range(n_clusters): + grid_membership = membership_matrix[nearest_indices, i].reshape(TIME.shape) + contour = ax.contourf(TIME, FREQ, grid_membership, levels=20, + cmap='viridis', alpha=options.alpha/2) + + # Overlay original points + max_membership_cluster = np.argmax(membership_matrix, axis=1) + colors = colormaps.get_cmap(options.colormap)( + np.linspace(0, 1, n_clusters))[max_membership_cluster] + + ax.scatter(time, frequencies, c=colors, s=options.marker_size**2, + alpha=options.alpha, edgecolors='black', linewidths=0.5) + + cbar = plt.colorbar(contour, ax=ax) + cbar.set_label('Membership Probability', fontsize=options.label_fontsize) + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Frequency', fontsize=options.label_fontsize) + ax.set_title(f'Soft Clustering ({method})', fontsize=options.title_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def save_clustering_plot( + fig: Figure, + filename: str, + options: Optional[ClusterPlotOptions] = None +) -> None: + """Save a clustering plot to file. + + Args: + fig: Figure to save + filename: Output filename + options: Plot configuration options + """ + if options is None: + options = ClusterPlotOptions() + + if not filename.endswith(f'.{options.save_format}'): + filename += f'.{options.save_format}' + + fig.savefig(filename, format=options.save_format, dpi=options.dpi, + bbox_inches='tight') + print(f"Saved clustering plot to: {filename}") \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/modal_plots.py b/src/tokeye/eigspec/vis/modal_plots.py new file mode 100644 index 0000000..a152939 --- /dev/null +++ b/src/tokeye/eigspec/vis/modal_plots.py @@ -0,0 +1,465 @@ +""" +Modal analysis visualization functions for eigspec package. + +This module provides specialized plotting functions for modal analysis results: +- Mode shape visualization with spatial coordinates +- MAC (Modal Assurance Criterion) matrix plots +- Frequency-damping stability diagrams +- Modal parameter evolution and validation plots + +Based on the MATLAB eigspec toolbox modal visualization functions: +- view_pcaspec_cluster_pattern.m - Modal pattern visualization +- mnfit_ptref.m - Point-reference mode shape fitting and display +- mnfit_clus_medoid.m - Cluster-based modal analysis plots +- Various modal plotting utilities in view_pcaspec_results.m +- Gaussian Process Regression plots in gp2dp.m +""" + +from typing import Optional, List, Tuple, Union, Dict, Any +import numpy as np +import numpy.typing as npt +import matplotlib.pyplot as plt +from matplotlib import colormaps +import matplotlib.colors as mcolors +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from dataclasses import dataclass +from scipy.interpolate import griddata + + +@dataclass +class ModalPlotOptions: + """Configuration options for modal analysis plots. + + Attributes: + figsize: Figure size (width, height) in inches + dpi: Figure DPI for resolution + fontsize: Base font size for text elements + title_fontsize: Font size for plot titles + label_fontsize: Font size for axis labels + legend_fontsize: Font size for legend text + grid: Whether to show grid lines + colormap: Colormap name for mode shape plots + contour_levels: Number of contour levels + marker_size: Size of scatter plot markers + line_width: Width of plot lines + alpha: Transparency level (0-1) + save_format: Default format for saving figures + """ + figsize: Tuple[float, float] = (10, 8) + dpi: int = 100 + fontsize: int = 12 + title_fontsize: int = 14 + label_fontsize: int = 12 + legend_fontsize: int = 10 + grid: bool = True + colormap: str = 'RdBu' + contour_levels: int = 20 + marker_size: float = 6.0 + line_width: float = 1.5 + alpha: float = 0.8 + save_format: str = 'png' + + +def plot_array_geometry( + coordinates: npt.NDArray[np.floating], + sensor_labels: Optional[List[str]] = None, + highlight_subsets: Optional[Dict[str, List[int]]] = None, + options: Optional[ModalPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot sensor array geometry. + + Args: + coordinates: Sensor coordinates (n_sensors x 2) - [theta, phi] or [x, y] + sensor_labels: Optional labels for each sensor + highlight_subsets: Dictionary of subset names to sensor indices + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ModalPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Plot all sensors + ax.scatter(coordinates[:, 0], coordinates[:, 1], + c='black', marker='x', s=options.marker_size**2, + alpha=options.alpha, label='All sensors') + + # Highlight subsets if provided + if highlight_subsets is not None: + colors = colormaps.get_cmap('tab10')(np.linspace(0, 1, len(highlight_subsets))) + for i, (subset_name, indices) in enumerate(highlight_subsets.items()): + ax.scatter(coordinates[indices, 0], coordinates[indices, 1], + c=[colors[i]], s=options.marker_size**2, + alpha=options.alpha, label=subset_name) + + # Add sensor labels if provided + if sensor_labels is not None: + for i, label in enumerate(sensor_labels): + ax.annotate(label, (coordinates[i, 0], coordinates[i, 1]), + xytext=(5, 5), textcoords='offset points', + fontsize=options.fontsize-2) + + ax.set_xlabel('θ [rad]' if np.max(coordinates[:, 0]) <= 2*np.pi else 'X', + fontsize=options.label_fontsize) + ax.set_ylabel('φ [rad]' if np.max(coordinates[:, 1]) <= 2*np.pi else 'Y', + fontsize=options.label_fontsize) + ax.set_title('Sensor Array Geometry', fontsize=options.title_fontsize) + + if highlight_subsets is not None: + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_shape_vectors( + coordinates: npt.NDArray[np.floating], + shape_vector: npt.NDArray[np.complex128], + mode_info: Optional[Dict[str, Any]] = None, + show_phase: bool = True, + normalize: bool = True, + options: Optional[ModalPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot mode shape vectors as scatter plot. + + Args: + coordinates: Sensor coordinates (n_sensors x 2) + shape_vector: Complex mode shape vector + mode_info: Optional dictionary with mode information (frequency, etc.) + show_phase: Whether to show phase information + normalize: Whether to normalize the shape vector + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ModalPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Normalize if requested + if normalize: + shape_vector = shape_vector / np.sqrt(np.sum(np.abs(shape_vector)**2)) + + # Extract real and imaginary parts + real_part = np.real(shape_vector) + imag_part = np.imag(shape_vector) + + # Plot real and imaginary parts + ax.scatter(coordinates[:, 0], real_part, + c='blue', marker='o', s=options.marker_size**2, + alpha=options.alpha, label='Real part') + ax.scatter(coordinates[:, 1], imag_part, + c='red', marker='s', s=options.marker_size**2, + alpha=options.alpha, label='Imaginary part') + + # Add phase information if requested + if show_phase: + phases = np.angle(shape_vector) + scatter = ax.scatter(coordinates[:, 0], coordinates[:, 1], + c=phases, cmap='hsv', s=options.marker_size**2, + alpha=options.alpha, marker='^') + cbar = plt.colorbar(scatter, ax=ax) + cbar.set_label('Phase [rad]', fontsize=options.label_fontsize) + + # Create title with mode information + title = 'Mode Shape Vector' + if mode_info is not None: + if 'frequency' in mode_info: + title += f" (f = {mode_info['frequency']:.2f})" + if 'mode_number' in mode_info: + title += f" - Mode {mode_info['mode_number']}" + + ax.set_xlabel('Sensor Coordinate', fontsize=options.label_fontsize) + ax.set_ylabel('Amplitude', fontsize=options.label_fontsize) + ax.set_title(title, fontsize=options.title_fontsize) + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_mode_shape_2d( + coordinates: npt.NDArray[np.floating], + shape_vector: npt.NDArray[np.complex128], + grid_size: Tuple[int, int] = (100, 75), + component: str = 'magnitude', + mode_info: Optional[Dict[str, Any]] = None, + options: Optional[ModalPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot 2D mode shape using interpolation and contours. + + Args: + coordinates: Sensor coordinates (n_sensors x 2) - [theta, phi] + shape_vector: Complex mode shape vector + grid_size: Size of interpolation grid (ntheta, nphi) + component: Component to plot ('magnitude', 'real', 'imag', 'phase') + mode_info: Optional dictionary with mode information + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ModalPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Create interpolation grid + theta_min, theta_max = coordinates[:, 0].min(), coordinates[:, 0].max() + phi_min, phi_max = coordinates[:, 1].min(), coordinates[:, 1].max() + + theta_grid = np.linspace(theta_min, theta_max, grid_size[0]) + phi_grid = np.linspace(phi_min, phi_max, grid_size[1]) + THETA, PHI = np.meshgrid(theta_grid, phi_grid) + + # Select component to plot + if component == 'magnitude': + values = np.abs(shape_vector) + cmap = options.colormap + label = 'Magnitude' + elif component == 'real': + values = np.real(shape_vector) + cmap = 'RdBu' + label = 'Real Part' + elif component == 'imag': + values = np.imag(shape_vector) + cmap = 'RdBu' + label = 'Imaginary Part' + elif component == 'phase': + values = np.angle(shape_vector) + cmap = 'hsv' + label = 'Phase [rad]' + else: + raise ValueError(f"Unknown component: {component}") + + # Interpolate to grid + grid_values = griddata(coordinates, values, (THETA, PHI), method='cubic') + + # Create contour plot + contour = ax.contourf(THETA, PHI, grid_values, levels=options.contour_levels, + cmap=cmap, alpha=options.alpha) + + # Add contour lines + ax.contour(THETA, PHI, grid_values, levels=options.contour_levels, + colors='black', alpha=0.3, linewidths=0.5) + + # Add sensor positions + ax.scatter(coordinates[:, 0], coordinates[:, 1], + c='black', marker='x', s=options.marker_size**2, alpha=1.0) + + # Add colorbar + cbar = plt.colorbar(contour, ax=ax) + cbar.set_label(label, fontsize=options.label_fontsize) + + # Create title + title = f'2D Mode Shape - {label}' + if mode_info is not None: + if 'frequency' in mode_info: + title += f" (f = {mode_info['frequency']:.2f})" + if 'mode_number' in mode_info: + title += f" - Mode {mode_info['mode_number']}" + + ax.set_xlabel('θ [rad]', fontsize=options.label_fontsize) + ax.set_ylabel('φ [rad]', fontsize=options.label_fontsize) + ax.set_title(title, fontsize=options.title_fontsize) + + return fig, ax + + +def plot_mode_shape_polar( + coordinates: npt.NDArray[np.floating], + shape_vector: npt.NDArray[np.complex128], + subset_indices: Optional[List[int]] = None, + smooth: bool = True, + mode_info: Optional[Dict[str, Any]] = None, + options: Optional[ModalPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot mode shape on a polar/toroidal array. + + Args: + coordinates: Sensor coordinates (n_sensors x 2) - [theta, phi] + shape_vector: Complex mode shape vector + subset_indices: Indices of sensors to use for polar plot + smooth: Whether to apply smoothing interpolation + mode_info: Optional dictionary with mode information + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = ModalPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + # Select subset if provided + if subset_indices is not None: + plot_coords = coordinates[subset_indices] + plot_shape = shape_vector[subset_indices] + else: + plot_coords = coordinates + plot_shape = shape_vector + + # Extract angles (assuming second coordinate is the angular one) + angles = plot_coords[:, 1] + + # Sort by angle for proper plotting + sort_indices = np.argsort(angles) + angles = angles[sort_indices] + plot_shape = plot_shape[sort_indices] + + # Create smooth interpolation if requested + if smooth: + from scipy.interpolate import interp1d + phi_smooth = np.linspace(0, 2*np.pi, 128) + + # Handle periodicity + extended_angles = np.concatenate([angles - 2*np.pi, angles, angles + 2*np.pi]) + extended_real = np.concatenate([np.real(plot_shape), np.real(plot_shape), + np.real(plot_shape)]) + extended_imag = np.concatenate([np.imag(plot_shape), np.imag(plot_shape), + np.imag(plot_shape)]) + + interp_real = interp1d(extended_angles, extended_real, kind='cubic') + interp_imag = interp1d(extended_angles, extended_imag, kind='cubic') + + smooth_real = interp_real(phi_smooth) + smooth_imag = interp_imag(phi_smooth) + + ax.plot(phi_smooth, smooth_real, 'b-', linewidth=options.line_width, + alpha=options.alpha, label='Real part') + ax.plot(phi_smooth, smooth_imag, 'r-', linewidth=options.line_width, + alpha=options.alpha, label='Imaginary part') + + # Plot original data points + ax.scatter(angles, np.real(plot_shape), c='blue', marker='o', + s=options.marker_size**2, alpha=options.alpha, + label='Real (data)' if smooth else 'Real part') + ax.scatter(angles, np.imag(plot_shape), c='red', marker='s', + s=options.marker_size**2, alpha=options.alpha, + label='Imag (data)' if smooth else 'Imaginary part') + + # Create title + title = 'Polar Mode Shape' + if mode_info is not None: + if 'frequency' in mode_info: + title += f" (f = {mode_info['frequency']:.2f})" + if 'mode_number' in mode_info: + title += f" - Mode {mode_info['mode_number']}" + + ax.set_xlabel('φ [rad]', fontsize=options.label_fontsize) + ax.set_ylabel('Amplitude', fontsize=options.label_fontsize) + ax.set_title(title, fontsize=options.title_fontsize) + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_mode_shapes( + coordinates: npt.NDArray[np.floating], + shape_vectors: npt.NDArray[np.complex128], + mode_info: Optional[List[Dict[str, Any]]] = None, + plot_types: List[str] = ['2d', 'polar'], + options: Optional[ModalPlotOptions] = None +) -> Figure: + """Create a comprehensive plot of multiple mode shapes. + + Args: + coordinates: Sensor coordinates (n_sensors x 2) + shape_vectors: Complex mode shape vectors (n_sensors x n_modes) + mode_info: List of dictionaries with mode information + plot_types: Types of plots to create ('2d', 'polar', 'vectors') + options: Plot configuration options + + Returns: + Figure object with subplots + """ + if options is None: + options = ModalPlotOptions() + + n_modes = shape_vectors.shape[1] + n_plots = len(plot_types) + + fig, axes = plt.subplots(n_modes, n_plots, figsize=(options.figsize[0]*n_plots, + options.figsize[1]*n_modes), + dpi=options.dpi) + + # Handle case of single mode or single plot type + if n_modes == 1: + axes = axes.reshape(1, -1) + if n_plots == 1: + axes = axes.reshape(-1, 1) + + for i in range(n_modes): + shape_vector = shape_vectors[:, i] + info = mode_info[i] if mode_info is not None else None + + for j, plot_type in enumerate(plot_types): + if plot_type == '2d': + plot_mode_shape_2d(coordinates, shape_vector, mode_info=info, + options=options, ax=axes[i, j]) + elif plot_type == 'polar': + plot_mode_shape_polar(coordinates, shape_vector, mode_info=info, + options=options, ax=axes[i, j]) + elif plot_type == 'vectors': + plot_shape_vectors(coordinates, shape_vector, mode_info=info, + options=options, ax=axes[i, j]) + + plt.tight_layout() + return fig + + +def save_modal_plot( + fig: Figure, + filename: str, + options: Optional[ModalPlotOptions] = None +) -> None: + """Save a modal plot to file. + + Args: + fig: Figure to save + filename: Output filename + options: Plot configuration options + """ + if options is None: + options = ModalPlotOptions() + + if not filename.endswith(f'.{options.save_format}'): + filename += f'.{options.save_format}' + + fig.savefig(filename, format=options.save_format, dpi=options.dpi, + bbox_inches='tight') + print(f"Saved modal plot to: {filename}") \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/spectral_plots.py b/src/tokeye/eigspec/vis/spectral_plots.py new file mode 100644 index 0000000..a40dc79 --- /dev/null +++ b/src/tokeye/eigspec/vis/spectral_plots.py @@ -0,0 +1,709 @@ +""" +Spectral plotting functions for eigspec package. + +This module provides specialized plotting functions for frequency-domain analysis: +- Power spectral density plots with modal overlays +- Frequency-time spectrograms and waterfall plots +- Coherence and phase plots for multi-channel analysis +- Eigenvalue and stability plots + +Based on the MATLAB eigspec toolbox spectral plotting functions: +- Various plotting routines in view_pcaspec_results.m +- Spectral analysis plots in eigspec_mmain.m +- FFT and frequency domain plotting utilities +- Coherence plotting in oddevenupdate_emc.m +""" + +from typing import Optional, List, Tuple, Union, Literal +import numpy as np +import numpy.typing as npt +import matplotlib.pyplot as plt +from matplotlib import colormaps +import matplotlib.colors as mcolors +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from dataclasses import dataclass, field + + +@dataclass +class SpectralPlotOptions: + """Configuration options for spectral analysis plots. + + Attributes: + figsize: Figure size (width, height) in inches + dpi: Figure DPI for resolution + fontsize: Base font size for text elements + title_fontsize: Font size for plot titles + label_fontsize: Font size for axis labels + legend_fontsize: Font size for legend text + grid: Whether to show grid lines + colormap: Colormap name for multi-series plots + marker_size: Size of scatter plot markers + line_width: Width of plot lines + alpha: Transparency level (0-1) + save_format: Default format for saving figures + """ + figsize: Tuple[float, float] = (12, 8) + dpi: int = 100 + fontsize: int = 12 + title_fontsize: int = 14 + label_fontsize: int = 12 + legend_fontsize: int = 10 + grid: bool = True + colormap: str = 'viridis' + marker_size: float = 4.0 + line_width: float = 1.5 + alpha: float = 0.8 + save_format: str = 'png' + + +def setup_spectral_figure( + nrows: int = 1, + ncols: int = 1, + options: Optional[SpectralPlotOptions] = None +) -> Tuple[Figure, Union[Axes, npt.NDArray]]: + """Set up a figure for spectral analysis plots. + + Args: + nrows: Number of subplot rows + ncols: Number of subplot columns + options: Plot configuration options + + Returns: + Figure and axes objects + """ + if options is None: + options = SpectralPlotOptions() + + fig, axes = plt.subplots(nrows, ncols, figsize=options.figsize, dpi=options.dpi) + + # Set default font sizes + plt.rcParams.update({ + 'font.size': options.fontsize, + 'axes.titlesize': options.title_fontsize, + 'axes.labelsize': options.label_fontsize, + 'legend.fontsize': options.legend_fontsize, + 'lines.linewidth': options.line_width, + 'lines.markersize': options.marker_size, + }) + + return fig, axes + + +def setup_single_spectral_figure( + options: Optional[SpectralPlotOptions] = None +) -> Tuple[Figure, Axes]: + """Set up a figure with a single axes for spectral analysis plots. + + Args: + options: Plot configuration options + + Returns: + Figure and single axes objects + """ + if options is None: + options = SpectralPlotOptions() + + fig, ax = plt.subplots(1, 1, figsize=options.figsize, dpi=options.dpi) + + # Set default font sizes + plt.rcParams.update({ + 'font.size': options.fontsize, + 'axes.titlesize': options.title_fontsize, + 'axes.labelsize': options.label_fontsize, + 'legend.fontsize': options.legend_fontsize, + 'lines.linewidth': options.line_width, + 'lines.markersize': options.marker_size, + }) + + return fig, ax + + +def plot_eigenvalue_evolution( + time: npt.NDArray[np.floating], + eigenvalues: npt.NDArray[np.floating], + n_retained: Optional[int] = None, + log_scale: bool = True, + highlight_time: Optional[float] = None, + options: Optional[SpectralPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot eigenvalue evolution over time. + + Args: + time: Time vector + eigenvalues: Eigenvalue matrix (time x eigenvalue_index) + n_retained: Number of retained eigenvalues to highlight + log_scale: Whether to use logarithmic scale for eigenvalues + highlight_time: Time point to highlight with vertical line + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = SpectralPlotOptions() + + if ax is None: + fig, ax = setup_single_spectral_figure(options=options) + else: + fig = ax.figure + + # Convert to log scale if requested + if log_scale: + plot_eigenvalues = np.log10(eigenvalues + 1e-16) # Add small value to avoid log(0) + ylabel = r'$\log_{10}(\lambda)$' + else: + plot_eigenvalues = eigenvalues + ylabel = r'$\lambda$' + + # Plot eigenvalues + n_eigs = eigenvalues.shape[1] + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_eigs)) + + for i in range(n_eigs): + if n_retained is not None and i < n_retained: + ax.plot(time, plot_eigenvalues[:, i], color=colors[i], + alpha=options.alpha, linewidth=options.line_width) + else: + ax.plot(time, plot_eigenvalues[:, i], 'k-', + alpha=options.alpha * 0.5, linewidth=options.line_width * 0.7) + + # Highlight specific time if provided + if highlight_time is not None: + ax.axvline(highlight_time, color='red', linestyle='--', linewidth=2) + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel(ylabel, fontsize=options.label_fontsize) + ax.set_title('Eigenvalue Evolution', fontsize=options.title_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_frequency_time( + time: npt.NDArray[np.floating], + frequencies: npt.NDArray[np.floating], + amplitudes: Optional[npt.NDArray[np.floating]] = None, + labels: Optional[List[str]] = None, + sampling_frequency: Optional[float] = None, + options: Optional[SpectralPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot frequency vs time scatter plot. + + Args: + time: Time vector + frequencies: Frequency values + amplitudes: Optional amplitude values for color coding + labels: Optional labels for different mode types + sampling_frequency: Sampling frequency for normalization + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = SpectralPlotOptions() + + if ax is None: + fig, ax = setup_single_spectral_figure(options=options) + else: + fig = ax.figure + + # Convert frequencies to appropriate units + if sampling_frequency is not None and sampling_frequency > 0: + plot_frequencies = frequencies / (1e3) # Convert to kHz + freq_label = 'Frequency [kHz]' + else: + plot_frequencies = frequencies + freq_label = 'Frequency [rad/sample]' + + # Plot based on whether we have labels or amplitudes + if labels is not None: + unique_labels = list(set(labels)) + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, len(unique_labels))) + + for i, label in enumerate(unique_labels): + mask = np.array(labels) == label + if label == 'nolabel': + ax.scatter(time[mask], plot_frequencies[mask], + c='black', s=options.marker_size**2, alpha=options.alpha, + label=label) + else: + ax.scatter(time[mask], plot_frequencies[mask], + c=[colors[i]], s=options.marker_size**2, alpha=options.alpha, + label=label) + ax.legend(fontsize=options.legend_fontsize) + + elif amplitudes is not None: + scatter = ax.scatter(time, plot_frequencies, c=amplitudes, + s=options.marker_size**2, alpha=options.alpha, + cmap=options.colormap) + plt.colorbar(scatter, ax=ax, label='Amplitude') + + else: + ax.scatter(time, plot_frequencies, s=options.marker_size**2, + alpha=options.alpha, color='blue') + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel(freq_label, fontsize=options.label_fontsize) + ax.set_title('Frequency vs Time', fontsize=options.title_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_phase_time( + time: npt.NDArray[np.floating], + eigenvalues: npt.NDArray[np.complex128], + mode_indices: Optional[List[int]] = None, + options: Optional[SpectralPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot phase angles of eigenvalues vs time. + + Args: + time: Time vector + eigenvalues: Complex eigenvalues + mode_indices: Indices of modes to plot + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = SpectralPlotOptions() + + if ax is None: + fig, ax = setup_single_spectral_figure(options=options) + else: + fig = ax.figure + + # Extract phase angles + phases = np.angle(eigenvalues) + + if mode_indices is not None: + for i in mode_indices: + ax.scatter(time, phases[:, i], s=options.marker_size**2, + alpha=options.alpha, label=f'Mode {i+1}') + ax.legend(fontsize=options.legend_fontsize) + else: + # Plot all modes as black dots + for i in range(phases.shape[1]): + ax.scatter(time, phases[:, i], c='black', s=options.marker_size**2, + alpha=options.alpha) + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel(r'Phase [rad]', fontsize=options.label_fontsize) + ax.set_title('Phase Evolution', fontsize=options.title_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_rms_time( + time: npt.NDArray[np.floating], + rms_values: npt.NDArray[np.floating], + frequencies: Optional[npt.NDArray[np.floating]] = None, + plot_type: Literal["rms_dot", "rms_b"] = "rms_dot", + labels: Optional[List[str]] = None, + options: Optional[SpectralPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot RMS values vs time. + + Args: + time: Time vector + rms_values: RMS amplitude values + frequencies: Optional frequency values for RMS(B) calculation + plot_type: Type of RMS plot ("rms_dot" or "rms_b") + labels: Optional labels for different modes + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = SpectralPlotOptions() + + if ax is None: + fig, ax = setup_single_spectral_figure(options=options) + else: + fig = ax.figure + + # Prepare data based on plot type + if plot_type == "rms_b" and frequencies is not None: + plot_values = (1e3) * rms_values / (2 * np.pi * frequencies) + ylabel = 'RMS (B) [mT]' + title = 'RMS Magnetic Field vs Time' + else: + plot_values = rms_values + ylabel = 'RMS (Ḃ) [T/s]' + title = 'RMS Magnetic Field Derivative vs Time' + + # Plot with or without labels + if labels is not None: + unique_labels = list(set(labels)) + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, len(unique_labels))) + + for i, label in enumerate(unique_labels): + mask = np.array(labels) == label + if label == 'nolabel': + ax.scatter(time[mask], plot_values[mask], + c='black', s=options.marker_size**2, alpha=options.alpha, + label=label) + else: + ax.scatter(time[mask], plot_values[mask], + c=[colors[i]], s=options.marker_size**2, alpha=options.alpha, + label=label) + ax.legend(fontsize=options.legend_fontsize) + else: + ax.scatter(time, plot_values, s=options.marker_size**2, + alpha=options.alpha, color='blue') + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel(ylabel, fontsize=options.label_fontsize) + ax.set_title(title, fontsize=options.title_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_spectral_summary( + time: npt.NDArray[np.floating], + eigenvalues: npt.NDArray[np.floating], + frequencies: npt.NDArray[np.floating], + rms_values: npt.NDArray[np.floating], + n_retained: Optional[int] = None, + sampling_frequency: Optional[float] = None, + options: Optional[SpectralPlotOptions] = None +) -> Figure: + """Create a summary plot of spectral analysis results. + + Args: + time: Time vector + eigenvalues: Eigenvalue matrix + frequencies: Frequency values + rms_values: RMS values + n_retained: Number of retained eigenvalues + sampling_frequency: Sampling frequency + options: Plot configuration options + + Returns: + Figure object with subplots + """ + if options is None: + options = SpectralPlotOptions() + + fig, axes = setup_spectral_figure(nrows=2, ncols=2, options=options) + + # Eigenvalue evolution + plot_eigenvalue_evolution(time, eigenvalues, n_retained=n_retained, + ax=axes[0, 0], options=options) + + # Frequency vs time + plot_frequency_time(time, frequencies, sampling_frequency=sampling_frequency, + ax=axes[0, 1], options=options) + + # Phase evolution (if eigenvalues are complex) + if np.iscomplexobj(eigenvalues): + plot_phase_time(time, eigenvalues.astype(np.complex128), ax=axes[1, 0], + options=options) + else: + axes[1, 0].text(0.5, 0.5, 'No phase data\n(real eigenvalues)', + ha='center', va='center', transform=axes[1, 0].transAxes) + axes[1, 0].set_title('Phase Evolution') + + # RMS vs time + plot_rms_time(time, rms_values, ax=axes[1, 1], options=options) + + plt.tight_layout() + return fig + + +def save_spectral_plot( + fig: Figure, + filename: str, + options: Optional[SpectralPlotOptions] = None +) -> None: + """Save a spectral plot to file. + + Args: + fig: Figure to save + filename: Output filename + options: Plot configuration options + """ + if options is None: + options = SpectralPlotOptions() + + if not filename.endswith(f'.{options.save_format}'): + filename += f'.{options.save_format}' + + fig.savefig(filename, format=options.save_format, dpi=options.dpi, + bbox_inches='tight') + print(f"Saved plot to: {filename}") + + +# ============================================================================= +# Main view_pcaspec visualization function (MATLAB equivalent) +# ============================================================================= + +def view_pcaspec( + analysis_result, + sensor_coordinates: Optional[npt.NDArray[np.floating]] = None, + time_point: Optional[float] = None, + options: Optional[SpectralPlotOptions] = None +) -> List[Figure]: + """ + Main visualization function for spectral analysis results. + + Python equivalent of MATLAB view_pcaspec.m that provides comprehensive + visualization of spectral analysis results including: + - Eigenvalue evolution over time + - Frequency-time spectrograms + - Modal shape visualization + - PCA component analysis + + Args: + analysis_result: Analysis result structure from eigspec functions + sensor_coordinates: Optional sensor coordinate matrix (N_sensors, 2) + time_point: Optional specific time point to visualize + options: Plot configuration options + + Returns: + List of Figure objects created + """ + if options is None: + options = SpectralPlotOptions() + + figures = [] + + # Extract time vector and block results + if hasattr(analysis_result, 'L') and analysis_result.L: + # Block-based analysis results + TL = np.array([block.centre_t for block in analysis_result.L]) + NBlock = len(TL) + + # Plot eigenvalue evolution if available + if hasattr(analysis_result.L[0], 'D'): + M = len(analysis_result.L[0].D) + r = analysis_result.L[0].mrep.m0.shape.shape[0] if hasattr(analysis_result.L[0].mrep.m0, 'shape') else 10 + + DD = np.zeros((NBlock, M)) + for jj, block in enumerate(analysis_result.L): + DD[jj, :] = block.D[:M] + + # Create eigenvalue evolution plot + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + + for mode_idx in range(min(r, M)): + ax.semilogy(TL, DD[:, mode_idx], 'o-', + markersize=options.marker_size, + linewidth=options.line_width, + label=f'Mode {mode_idx+1}') + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Eigenvalue Magnitude', fontsize=options.label_fontsize) + ax.set_title('PCA Eigenvalue Evolution', fontsize=options.title_fontsize) + ax.grid(options.grid) + ax.legend(fontsize=options.legend_fontsize) + + figures.append(fig) + + # Plot frequency evolution if available + if hasattr(analysis_result.L[0].mrep, 'm0') and hasattr(analysis_result.L[0].mrep.m0, 'lambda'): + n_modes_max = max(len(block.mrep.imode) for block in analysis_result.L if hasattr(block.mrep, 'imode')) + + if n_modes_max > 0: + freq_data = np.full((NBlock, n_modes_max), np.nan) + + for jj, block in enumerate(analysis_result.L): + if hasattr(block.mrep, 'imode') and hasattr(block.mrep.m0, 'lambda'): + # 'lambda' is a keyword: attribute access needs getattr + eigenvalues = getattr(block.mrep.m0, 'lambda') + for idx, mode_idx in enumerate(block.mrep.imode): + if idx < n_modes_max and mode_idx < len(eigenvalues): + eigenval = eigenvalues[mode_idx] + # Convert complex eigenvalue to frequency + if np.iscomplexobj(eigenval): + freq_data[jj, idx] = np.abs(np.angle(eigenval)) + else: + freq_data[jj, idx] = np.abs(eigenval) + + # Create frequency-time plot + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + + for mode_idx in range(n_modes_max): + valid_mask = ~np.isnan(freq_data[:, mode_idx]) + if np.any(valid_mask): + ax.plot(TL[valid_mask], freq_data[valid_mask, mode_idx], 'o-', + markersize=options.marker_size, + linewidth=options.line_width, + label=f'Mode {mode_idx+1}') + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Frequency (rad/sample)', fontsize=options.label_fontsize) + ax.set_title('Modal Frequency Evolution', fontsize=options.title_fontsize) + ax.grid(options.grid) + ax.legend(fontsize=options.legend_fontsize) + + figures.append(fig) + + # Single time point visualization + if time_point is not None and hasattr(analysis_result, 'L'): + if time_point < TL.min() or time_point > TL.max(): + print(f"Warning: time point {time_point} out of range [{TL.min():.3f}, {TL.max():.3f}]") + return figures + + # Find closest time block + ll = np.argmin(np.abs(TL - time_point)) + Lll = analysis_result.L[ll] + + # Plot mode shapes if sensor coordinates provided + if sensor_coordinates is not None and hasattr(Lll.mrep, 'imode'): + for mm, mode_idx in enumerate(Lll.mrep.imode): + if hasattr(Lll.mrep.m0, 'shape') and mode_idx < len(Lll.mrep.m0.shape): + fig, axes = plt.subplots(1, 2, figsize=(2*options.figsize[0]/3, options.figsize[1]), + dpi=options.dpi) + + mode_shape = Lll.mrep.m0.shape[mode_idx] + freq = getattr(Lll.mrep.m0, 'lambda')[mode_idx] if hasattr(Lll.mrep.m0, 'lambda') else 0 + + # Normalize mode shape + mode_shape = mode_shape / np.sqrt(np.vdot(mode_shape, mode_shape)) + + # Real part + axes[0].scatter(sensor_coordinates[:, 0], sensor_coordinates[:, 1], + c=np.real(mode_shape), s=options.marker_size*20, + cmap=options.colormap, alpha=options.alpha) + axes[0].set_title(f'Mode {mm+1} - Real Part', fontsize=options.title_fontsize) + axes[0].set_xlabel('X Coordinate', fontsize=options.label_fontsize) + axes[0].set_ylabel('Y Coordinate', fontsize=options.label_fontsize) + axes[0].grid(options.grid) + + # Imaginary part + axes[1].scatter(sensor_coordinates[:, 0], sensor_coordinates[:, 1], + c=np.imag(mode_shape), s=options.marker_size*20, + cmap=options.colormap, alpha=options.alpha) + axes[1].set_title(f'Mode {mm+1} - Imaginary Part', fontsize=options.title_fontsize) + axes[1].set_xlabel('X Coordinate', fontsize=options.label_fontsize) + axes[1].set_ylabel('Y Coordinate', fontsize=options.label_fontsize) + axes[1].grid(options.grid) + + fig.suptitle(f'Mode Shape at t={time_point:.3f}s, f={np.abs(freq):.3f}', + fontsize=options.title_fontsize) + plt.tight_layout() + + figures.append(fig) + + # Summary statistics plot + if hasattr(analysis_result, 'L') and analysis_result.L: + fig, axes = plt.subplots(2, 2, figsize=options.figsize, dpi=options.dpi) + axes = axes.flatten() + + # Number of modes per time block + n_modes_per_block = [len(block.mrep.imode) if hasattr(block.mrep, 'imode') else 0 + for block in analysis_result.L] + + axes[0].plot(TL, n_modes_per_block, 'o-', + markersize=options.marker_size, linewidth=options.line_width) + axes[0].set_xlabel('Time', fontsize=options.label_fontsize) + axes[0].set_ylabel('Number of Modes', fontsize=options.label_fontsize) + axes[0].set_title('Mode Count Evolution', fontsize=options.title_fontsize) + axes[0].grid(options.grid) + + # Processing time per block if available + if hasattr(analysis_result.L[0], 'block_processing_time'): + proc_times = [block.block_processing_time for block in analysis_result.L] + axes[1].plot(TL, proc_times, 'o-', + markersize=options.marker_size, linewidth=options.line_width) + axes[1].set_xlabel('Time', fontsize=options.label_fontsize) + axes[1].set_ylabel('Processing Time (s)', fontsize=options.label_fontsize) + axes[1].set_title('Block Processing Time', fontsize=options.title_fontsize) + axes[1].grid(options.grid) + else: + axes[1].text(0.5, 0.5, 'No timing data\navailable', + ha='center', va='center', transform=axes[1].transAxes) + axes[1].set_title('Block Processing Time', fontsize=options.title_fontsize) + + # Mode frequency distribution + all_freqs = [] + for block in analysis_result.L: + if hasattr(block.mrep, 'm0') and hasattr(block.mrep.m0, 'lambda'): + # 'lambda' is a keyword: attribute access needs getattr + eigenvalues = getattr(block.mrep.m0, 'lambda') + for mode_idx in block.mrep.imode: + if mode_idx < len(eigenvalues): + eigenval = eigenvalues[mode_idx] + freq = np.abs(np.angle(eigenval)) if np.iscomplexobj(eigenval) else np.abs(eigenval) + all_freqs.append(freq) + + if all_freqs: + axes[2].hist(all_freqs, bins=20, alpha=options.alpha, edgecolor='black') + axes[2].set_xlabel('Frequency (rad/sample)', fontsize=options.label_fontsize) + axes[2].set_ylabel('Count', fontsize=options.label_fontsize) + axes[2].set_title('Mode Frequency Distribution', fontsize=options.title_fontsize) + axes[2].grid(options.grid) + else: + axes[2].text(0.5, 0.5, 'No frequency data\navailable', + ha='center', va='center', transform=axes[2].transAxes) + axes[2].set_title('Mode Frequency Distribution', fontsize=options.title_fontsize) + + # Hide unused subplot + axes[3].axis('off') + + plt.tight_layout() + figures.append(fig) + + return figures + + +def view_pcaspec_results( + analysis_result, + sensor_coordinates: npt.NDArray[np.floating], + threshold: float = 0.1, + options: Optional[SpectralPlotOptions] = None +) -> List[Figure]: + """ + Results visualization with minimal classification. + + Python equivalent of MATLAB view_pcaspec_results.m for displaying + analysis results with basic mode classification. + + Args: + analysis_result: Analysis result structure + sensor_coordinates: Sensor coordinate matrix (3, 2) for toroidal geometry + threshold: Classification threshold + options: Plot configuration options + + Returns: + List of Figure objects created + """ + if sensor_coordinates.shape != (3, 2): + raise ValueError("sensor_coordinates must be (3, 2) for toroidal geometry") + + if options is None: + options = SpectralPlotOptions() + + figures = [] + + # Extract results similar to MATLAB version + if hasattr(analysis_result, 'L') and analysis_result.L: + # Create frequency evolution plot + fig = view_pcaspec(analysis_result, sensor_coordinates, options=options) + figures.extend(fig) + + return figures \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/time_series_plots.py b/src/tokeye/eigspec/vis/time_series_plots.py new file mode 100644 index 0000000..4918f30 --- /dev/null +++ b/src/tokeye/eigspec/vis/time_series_plots.py @@ -0,0 +1,376 @@ +""" +Time series visualization functions for eigspec package. + +This module provides specialized plotting functions for time-domain analysis: +- Multi-channel time series plots with modal overlays +- Block-wise analysis timeline visualization +- RMS evolution and amplitude tracking plots +- Time-frequency analysis and trend visualization + +Based on the MATLAB eigspec toolbox time series plotting functions: +- Time series plotting in view_pcaspec_results.m +- Block-wise timeline plots in eigspec_mmain.m +- RMS and amplitude evolution plots +- collect_rep_refvec_trace.m - Reference vector trace plotting +- Various time-domain plotting utilities +""" + +from typing import Optional, List, Tuple, Union, Dict, Any +import numpy as np +import numpy.typing as npt +import matplotlib.pyplot as plt +from matplotlib import colormaps +from matplotlib.figure import Figure +from matplotlib.axes import Axes +from dataclasses import dataclass + + +@dataclass +class TimeSeriesPlotOptions: + """Configuration options for time series plots. + + Attributes: + figsize: Figure size (width, height) in inches + dpi: Figure DPI for resolution + fontsize: Base font size for text elements + title_fontsize: Font size for plot titles + label_fontsize: Font size for axis labels + legend_fontsize: Font size for legend text + grid: Whether to show grid lines + colormap: Colormap name for multi-series plots + marker_size: Size of scatter plot markers + line_width: Width of plot lines + alpha: Transparency level (0-1) + save_format: Default format for saving figures + """ + figsize: Tuple[float, float] = (14, 8) + dpi: int = 100 + fontsize: int = 12 + title_fontsize: int = 14 + label_fontsize: int = 12 + legend_fontsize: int = 10 + grid: bool = True + colormap: str = 'tab10' + marker_size: float = 4.0 + line_width: float = 1.5 + alpha: float = 0.8 + save_format: str = 'png' + + +def plot_time_traces( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + channel_labels: Optional[List[str]] = None, + plot_type: str = 'line', + normalize: bool = False, + options: Optional[TimeSeriesPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot time series traces for multiple channels. + + Args: + time: Time vector + data: Data matrix (n_time x n_channels) + channel_labels: Optional labels for channels + plot_type: Type of plot ('line', 'scatter', 'both') + normalize: Whether to normalize each trace + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = TimeSeriesPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + n_channels = data.shape[1] + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_channels)) + + # Normalize data if requested + if normalize: + plot_data = data / np.max(np.abs(data), axis=0) + else: + plot_data = data + + # Plot each channel + for i in range(n_channels): + label = channel_labels[i] if channel_labels else f'Channel {i+1}' + + if plot_type == 'line': + ax.plot(time, plot_data[:, i], color=colors[i], + linewidth=options.line_width, alpha=options.alpha, label=label) + elif plot_type == 'scatter': + ax.scatter(time, plot_data[:, i], c=[colors[i]], + s=options.marker_size**2, alpha=options.alpha, label=label) + elif plot_type == 'both': + ax.plot(time, plot_data[:, i], color=colors[i], + linewidth=options.line_width, alpha=options.alpha) + ax.scatter(time, plot_data[:, i], c=[colors[i]], + s=options.marker_size**2, alpha=options.alpha, label=label) + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Amplitude', fontsize=options.label_fontsize) + ax.set_title('Time Series Traces', fontsize=options.title_fontsize) + + if n_channels <= 10: + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def plot_prototype_traces( + time_data: List[npt.NDArray[np.floating]], + frequency_data: List[npt.NDArray[np.floating]], + rms_data: List[npt.NDArray[np.floating]], + mode_labels: List[str], + plot_frequency: bool = True, + plot_rms: bool = True, + filter_beta: Optional[float] = None, + options: Optional[TimeSeriesPlotOptions] = None +) -> Figure: + """Plot prototype trace data for multiple modes. + + Args: + time_data: List of time vectors for each mode + frequency_data: List of frequency data for each mode + rms_data: List of RMS data for each mode + mode_labels: Labels for each mode (e.g., 'm/n=1/1') + plot_frequency: Whether to plot frequency traces + plot_rms: Whether to plot RMS traces + filter_beta: Optional filtering parameter + options: Plot configuration options + + Returns: + Figure object with subplots + """ + if options is None: + options = TimeSeriesPlotOptions() + + n_plots = sum([plot_frequency, plot_rms]) + fig, axes = plt.subplots(n_plots, 1, figsize=(options.figsize[0], + options.figsize[1]*n_plots/2), + dpi=options.dpi) + + if n_plots == 1: + axes = [axes] + + n_modes = len(time_data) + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_modes)) + + plot_idx = 0 + + if plot_frequency: + ax = axes[plot_idx] + for i in range(n_modes): + ax.scatter(time_data[i] * 1e3, frequency_data[i] / 1e3, + c=[colors[i]], marker='o', s=options.marker_size**2, + alpha=options.alpha, label=mode_labels[i]) + + ax.set_xlabel('Time [ms]', fontsize=options.label_fontsize) + ax.set_ylabel('Frequency [kHz]', fontsize=options.label_fontsize) + ax.set_title('Prototype Frequency Traces', fontsize=options.title_fontsize) + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + plot_idx += 1 + + if plot_rms: + ax = axes[plot_idx] + for i in range(n_modes): + ax.scatter(time_data[i] * 1e3, rms_data[i], + c=[colors[i]], marker='o', s=options.marker_size**2, + alpha=options.alpha, label=mode_labels[i]) + + ax.set_xlabel('Time [ms]', fontsize=options.label_fontsize) + ax.set_ylabel('RMS [T/s]', fontsize=options.label_fontsize) + + title = 'Prototype RMS Traces' + if filter_beta is not None: + title += f' (filtered, β={filter_beta:.2f})' + ax.set_title(title, fontsize=options.title_fontsize) + + ax.legend(fontsize=options.legend_fontsize) + + if options.grid: + ax.grid(True, alpha=0.3) + + plt.tight_layout() + return fig + + +def plot_filtered_traces( + time: npt.NDArray[np.floating], + original_data: npt.NDArray[np.floating], + filtered_data: npt.NDArray[np.floating], + filter_params: Dict[str, Any], + channel_indices: Optional[List[int]] = None, + options: Optional[TimeSeriesPlotOptions] = None +) -> Figure: + """Plot comparison of original and filtered traces. + + Args: + time: Time vector + original_data: Original data matrix + filtered_data: Filtered data matrix + filter_params: Dictionary with filter parameters + channel_indices: Optional indices of channels to plot + options: Plot configuration options + + Returns: + Figure object with subplots + """ + if options is None: + options = TimeSeriesPlotOptions() + + if channel_indices is None: + channel_indices = list(range(min(4, original_data.shape[1]))) + + n_channels = len(channel_indices) + fig, axes = plt.subplots(n_channels, 1, figsize=(options.figsize[0], + options.figsize[1]*n_channels/3), + dpi=options.dpi) + + if n_channels == 1: + axes = [axes] + + for i, ch_idx in enumerate(channel_indices): + ax = axes[i] + + ax.plot(time, original_data[:, ch_idx], 'b-', + linewidth=options.line_width, alpha=options.alpha, + label='Original') + ax.plot(time, filtered_data[:, ch_idx], 'r-', + linewidth=options.line_width, alpha=options.alpha, + label='Filtered') + + ax.set_ylabel(f'Channel {ch_idx+1}', fontsize=options.label_fontsize) + ax.legend(fontsize=options.legend_fontsize-2) + + if options.grid: + ax.grid(True, alpha=0.3) + + axes[-1].set_xlabel('Time', fontsize=options.label_fontsize) + + # Create title with filter information + filter_info = ', '.join([f'{k}={v}' for k, v in filter_params.items()]) + fig.suptitle(f'Original vs Filtered Data ({filter_info})', + fontsize=options.title_fontsize) + + plt.tight_layout() + return fig + + +def plot_multi_channel_overlay( + time: npt.NDArray[np.floating], + data: npt.NDArray[np.floating], + overlay_style: str = 'offset', + channel_labels: Optional[List[str]] = None, + offset_scale: float = 1.0, + options: Optional[TimeSeriesPlotOptions] = None, + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Plot multiple channels with offset or overlay styling. + + Args: + time: Time vector + data: Data matrix (n_time x n_channels) + overlay_style: Style of overlay ('offset', 'transparent', 'normalized') + channel_labels: Optional labels for channels + offset_scale: Scaling factor for offset style + options: Plot configuration options + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if options is None: + options = TimeSeriesPlotOptions() + + if ax is None: + fig, ax = plt.subplots(figsize=options.figsize, dpi=options.dpi) + else: + fig = ax.figure + + n_channels = data.shape[1] + colors = colormaps.get_cmap(options.colormap)(np.linspace(0, 1, n_channels)) + + if overlay_style == 'offset': + # Calculate offset based on data range + data_range = np.max(data) - np.min(data) + offset = data_range * offset_scale + + for i in range(n_channels): + label = channel_labels[i] if channel_labels else f'Ch {i+1}' + offset_data = data[:, i] + i * offset + ax.plot(time, offset_data, color=colors[i], + linewidth=options.line_width, alpha=options.alpha, label=label) + + # Set custom y-tick labels + if channel_labels: + y_positions = [i * offset for i in range(n_channels)] + ax.set_yticks(y_positions) + ax.set_yticklabels(channel_labels) + + elif overlay_style == 'transparent': + for i in range(n_channels): + label = channel_labels[i] if channel_labels else f'Ch {i+1}' + ax.plot(time, data[:, i], color=colors[i], + linewidth=options.line_width, alpha=options.alpha*0.7, label=label) + + elif overlay_style == 'normalized': + # Normalize each channel to [0, 1] range + normalized_data = np.zeros_like(data) + for i in range(n_channels): + ch_data = data[:, i] + normalized_data[:, i] = (ch_data - np.min(ch_data)) / (np.max(ch_data) - np.min(ch_data)) + + for i in range(n_channels): + label = channel_labels[i] if channel_labels else f'Ch {i+1}' + ax.plot(time, normalized_data[:, i], color=colors[i], + linewidth=options.line_width, alpha=options.alpha, label=label) + + ax.set_xlabel('Time', fontsize=options.label_fontsize) + ax.set_ylabel('Amplitude', fontsize=options.label_fontsize) + ax.set_title(f'Multi-Channel Overlay ({overlay_style})', fontsize=options.title_fontsize) + + if n_channels <= 15: + ax.legend(fontsize=options.legend_fontsize, bbox_to_anchor=(1.05, 1), + loc='upper left') + + if options.grid: + ax.grid(True, alpha=0.3) + + return fig, ax + + +def save_time_series_plot( + fig: Figure, + filename: str, + options: Optional[TimeSeriesPlotOptions] = None +) -> None: + """Save a time series plot to file. + + Args: + fig: Figure to save + filename: Output filename + options: Plot configuration options + """ + if options is None: + options = TimeSeriesPlotOptions() + + if not filename.endswith(f'.{options.save_format}'): + filename += f'.{options.save_format}' + + fig.savefig(filename, format=options.save_format, dpi=options.dpi, + bbox_inches='tight') + print(f"Saved time series plot to: {filename}") \ No newline at end of file diff --git a/src/tokeye/eigspec/vis/utility_plots.py b/src/tokeye/eigspec/vis/utility_plots.py new file mode 100644 index 0000000..c4fa228 --- /dev/null +++ b/src/tokeye/eigspec/vis/utility_plots.py @@ -0,0 +1,363 @@ +""" +Utility plotting functions for eigspec package. + +This module provides general-purpose plotting utilities and diagnostic plots: +- Analysis diagnostics and validation plots +- Parameter convergence and optimization plots +- Error analysis and statistical visualization +- General-purpose scientific plotting utilities + +Based on the MATLAB eigspec toolbox utility plotting functions: +- Diagnostic plots throughout the MATLAB codebase +- GCV and cross-validation plots in gcv1dp.m and df1dp.m +- Optimization and convergence plots +- Statistical analysis and error plotting utilities +- General plotting utilities used across the toolbox +""" + +from typing import Optional, List, Tuple, Union, Dict, Any +import numpy as np +import numpy.typing as npt +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +from matplotlib.figure import Figure +from matplotlib.axes import Axes + + +def plot_quantiles( + data: npt.NDArray[np.floating], + quantiles: Optional[npt.NDArray[np.floating]] = None, + column_labels: Optional[List[str]] = None, + figsize: Tuple[float, float] = (10, 6), + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Create a quantile plot showing statistical distribution for each column. + + This function is equivalent to the MATLAB qplot function, showing + percentile ranges for each data column. + + Args: + data: Data matrix (n_samples x n_features) + quantiles: Quantile levels to plot (default: [0.01, 0.1, 0.33, 0.5, 0.67, 0.9, 0.99]) + column_labels: Labels for each column + figsize: Figure size + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if quantiles is None: + quantiles = np.array([0.01, 0.1, 1/3, 0.5, 2/3, 0.9, 0.99]) + + if len(quantiles) != 7: + raise ValueError("quantiles must have exactly 7 elements") + + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.figure + + n_samples, n_features = data.shape + + if column_labels is None: + column_labels = [f'#{i+1}' for i in range(n_features)] + elif len(column_labels) != n_features: + raise ValueError("Number of labels must match number of columns") + + # Sort quantiles to ensure proper ordering + quantiles = np.sort(quantiles) + + # Calculate quantile values for each column + for i in range(n_features): + col_data = np.sort(data[:, i]) + quantile_indices = np.round(quantiles * (n_samples - 1)).astype(int) + quantile_values = col_data[quantile_indices] + + y_pos = i + 1 + + # Plot different ranges with different line styles and colors + # Extreme range (1-99%) + ax.plot([quantile_values[0], quantile_values[6]], [y_pos, y_pos], + 'g-', linewidth=2, alpha=0.7) + + # Outer range (10-90%) + ax.plot([quantile_values[1], quantile_values[5]], [y_pos, y_pos], + 'b-', linewidth=3, alpha=0.8) + + # Inner range (33-67%) + ax.plot([quantile_values[2], quantile_values[4]], [y_pos, y_pos], + 'r-', linewidth=4, alpha=0.9) + + # Median line + ax.plot([quantile_values[3], quantile_values[3]], [y_pos-0.4, y_pos+0.4], + 'k-', linewidth=2) + + # Set y-axis properties + ax.set_ylim(0, n_features + 1) + ax.set_yticks(range(1, n_features + 1)) + ax.set_yticklabels(column_labels) + + ax.set_xlabel('Value') + ax.set_title('Quantile Plot') + ax.grid(True, alpha=0.3) + + # Add legend + from matplotlib.lines import Line2D + legend_elements = [ + Line2D([0], [0], color='g', linewidth=2, alpha=0.7, label='1-99%'), + Line2D([0], [0], color='b', linewidth=3, alpha=0.8, label='10-90%'), + Line2D([0], [0], color='r', linewidth=4, alpha=0.9, label='33-67%'), + Line2D([0], [0], color='k', linewidth=2, label='Median') + ] + ax.legend(handles=legend_elements, loc='upper right') + + return fig, ax + + +def plot_statistical_summary( + data: npt.NDArray[np.floating], + labels: Optional[List[str]] = None, + stats: List[str] = ['mean', 'std', 'min', 'max'], + figsize: Tuple[float, float] = (12, 8) +) -> Figure: + """Create a comprehensive statistical summary plot. + + Args: + data: Data matrix (n_samples x n_features) + labels: Labels for each feature + stats: List of statistics to compute ('mean', 'std', 'min', 'max', 'median') + figsize: Figure size + + Returns: + Figure object with subplots + """ + n_features = data.shape[1] + n_stats = len(stats) + + if labels is None: + labels = [f'Feature {i+1}' for i in range(n_features)] + + fig, axes = plt.subplots(2, 2, figsize=figsize) + axes = axes.flatten() + + x_pos = np.arange(n_features) + + # Calculate statistics + stat_values = {} + for stat in stats: + if stat == 'mean': + stat_values[stat] = np.mean(data, axis=0) + elif stat == 'std': + stat_values[stat] = np.std(data, axis=0) + elif stat == 'min': + stat_values[stat] = np.min(data, axis=0) + elif stat == 'max': + stat_values[stat] = np.max(data, axis=0) + elif stat == 'median': + stat_values[stat] = np.median(data, axis=0) + + # Plot each statistic + for i, stat in enumerate(stats[:4]): # Limit to 4 subplots + if i >= len(axes): + break + + axes[i].bar(x_pos, stat_values[stat], alpha=0.7, + color=plt.cm.tab10(i)) + axes[i].set_title(f'{stat.capitalize()}') + axes[i].set_xticks(x_pos) + axes[i].set_xticklabels(labels, rotation=45, ha='right') + axes[i].grid(True, alpha=0.3) + + # Hide unused subplots + for i in range(len(stats), len(axes)): + axes[i].set_visible(False) + + plt.tight_layout() + return fig + + +def create_colormap( + colors: List[str], + name: str = 'custom', + n_segments: int = 256 +) -> mcolors.LinearSegmentedColormap: + """Create a custom colormap from a list of colors. + + Args: + colors: List of color names or hex codes + name: Name for the colormap + n_segments: Number of segments in the colormap + + Returns: + Custom colormap object + """ + if len(colors) < 2: + raise ValueError("At least 2 colors are required") + + # Convert colors to RGB if needed + rgb_colors = [] + for color in colors: + if isinstance(color, str): + rgb_colors.append(mcolors.to_rgb(color)) + else: + rgb_colors.append(color) + + # Create colormap + cmap = mcolors.LinearSegmentedColormap.from_list(name, rgb_colors, N=n_segments) + return cmap + + +def setup_figure_style( + style: str = 'default', + font_scale: float = 1.0, + grid: bool = True +) -> None: + """Set up matplotlib figure style and parameters. + + Args: + style: Style name ('default', 'publication', 'presentation') + font_scale: Scale factor for font sizes + grid: Whether to show grids by default + """ + if style == 'publication': + # Publication-ready style + plt.rcParams.update({ + 'font.size': 10 * font_scale, + 'axes.titlesize': 12 * font_scale, + 'axes.labelsize': 10 * font_scale, + 'xtick.labelsize': 9 * font_scale, + 'ytick.labelsize': 9 * font_scale, + 'legend.fontsize': 9 * font_scale, + 'figure.titlesize': 14 * font_scale, + 'lines.linewidth': 1.0, + 'lines.markersize': 4, + 'axes.grid': grid, + 'grid.alpha': 0.3, + 'figure.dpi': 150, + 'savefig.dpi': 300, + 'savefig.bbox': 'tight', + 'figure.facecolor': 'white', + 'axes.facecolor': 'white', + }) + + elif style == 'presentation': + # Presentation style with larger fonts + plt.rcParams.update({ + 'font.size': 14 * font_scale, + 'axes.titlesize': 18 * font_scale, + 'axes.labelsize': 16 * font_scale, + 'xtick.labelsize': 12 * font_scale, + 'ytick.labelsize': 12 * font_scale, + 'legend.fontsize': 14 * font_scale, + 'figure.titlesize': 20 * font_scale, + 'lines.linewidth': 2.0, + 'lines.markersize': 6, + 'axes.grid': grid, + 'grid.alpha': 0.3, + 'figure.dpi': 100, + 'savefig.dpi': 150, + 'savefig.bbox': 'tight', + }) + + else: # default + # Default matplotlib style with minor adjustments + plt.rcParams.update({ + 'font.size': 12 * font_scale, + 'axes.titlesize': 14 * font_scale, + 'axes.labelsize': 12 * font_scale, + 'xtick.labelsize': 11 * font_scale, + 'ytick.labelsize': 11 * font_scale, + 'legend.fontsize': 11 * font_scale, + 'figure.titlesize': 16 * font_scale, + 'lines.linewidth': 1.5, + 'lines.markersize': 5, + 'axes.grid': grid, + 'grid.alpha': 0.3, + 'figure.dpi': 100, + 'savefig.dpi': 200, + 'savefig.bbox': 'tight', + }) + + +def create_comparison_plot( + data_dict: Dict[str, npt.NDArray[np.floating]], + x_values: Optional[npt.NDArray[np.floating]] = None, + plot_type: str = 'line', + title: str = 'Comparison Plot', + xlabel: str = 'X', + ylabel: str = 'Y', + figsize: Tuple[float, float] = (10, 6), + ax: Optional[Axes] = None +) -> Tuple[Figure, Axes]: + """Create a comparison plot for multiple datasets. + + Args: + data_dict: Dictionary of {label: data_array} pairs + x_values: X-axis values (uses indices if None) + plot_type: Type of plot ('line', 'scatter', 'bar') + title: Plot title + xlabel: X-axis label + ylabel: Y-axis label + figsize: Figure size + ax: Existing axes to plot on + + Returns: + Figure and axes objects + """ + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.figure + + n_datasets = len(data_dict) + colors = plt.cm.tab10(np.linspace(0, 1, n_datasets)) + + for i, (label, data) in enumerate(data_dict.items()): + if x_values is None: + x = np.arange(len(data)) + else: + x = x_values[:len(data)] + + if plot_type == 'line': + ax.plot(x, data, color=colors[i], linewidth=2, alpha=0.8, label=label) + elif plot_type == 'scatter': + ax.scatter(x, data, color=colors[i], s=40, alpha=0.8, label=label) + elif plot_type == 'bar': + width = 0.8 / n_datasets + offset = (i - n_datasets/2 + 0.5) * width + ax.bar(x + offset, data, width=width, color=colors[i], + alpha=0.8, label=label) + + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend() + ax.grid(True, alpha=0.3) + + return fig, ax + + +def save_all_figures( + figures: List[Figure], + base_filename: str, + formats: List[str] = ['png'], + dpi: int = 200, + close_after_save: bool = True +) -> None: + """Save multiple figures with sequential numbering. + + Args: + figures: List of figure objects to save + base_filename: Base filename (will add numbers and extensions) + formats: List of formats to save ('png', 'pdf', 'svg', etc.) + dpi: DPI for raster formats + close_after_save: Whether to close figures after saving + """ + for i, fig in enumerate(figures): + for fmt in formats: + filename = f"{base_filename}_{i+1:02d}.{fmt}" + fig.savefig(filename, format=fmt, dpi=dpi, bbox_inches='tight') + print(f"Saved: {filename}") + + if close_after_save: + plt.close(fig) \ No newline at end of file diff --git a/src/tokeye/elmspec/__init__.py b/src/tokeye/elmspec/__init__.py new file mode 100644 index 0000000..e307f6c --- /dev/null +++ b/src/tokeye/elmspec/__init__.py @@ -0,0 +1,27 @@ +"""ELM detection from TokEye's transient-activity channel. + +``tokeye elmspec`` runs the segmentation model on spectrograms and turns the +transient channel (``mask[1]``) into discrete ELM events: time intervals, +counts, and ELM frequency. Pure-numpy event extraction lives in +:mod:`tokeye.elmspec.events`; the model plumbing is in the CLI handler. +""" + +from __future__ import annotations + +from tokeye.elmspec.events import ( + ElmEvent, + column_activity, + extract_elm_events, + summarize, + write_events_csv, + write_summary_csv, +) + +__all__ = [ + "ElmEvent", + "column_activity", + "extract_elm_events", + "summarize", + "write_events_csv", + "write_summary_csv", +] diff --git a/src/tokeye/elmspec/events.py b/src/tokeye/elmspec/events.py new file mode 100644 index 0000000..6d84b31 --- /dev/null +++ b/src/tokeye/elmspec/events.py @@ -0,0 +1,155 @@ +"""Pure-numpy ELM event extraction from a transient-activity mask. + +Input is the transient channel of a TokEye mask (``mask[1]``, shape ``(H, W)``, +values in [0, 1]). An ELM shows up as a broadband vertical stripe: many +frequency bins active in the same time column. Detection is therefore +column-wise: threshold the mask, measure the active fraction per column, +mark columns above ``activity_min``, close small gaps, and report the +remaining contiguous runs as events. +""" + +from __future__ import annotations + +import csv +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from pathlib import Path + + +@dataclass(frozen=True) +class ElmEvent: + start_col: int + end_col: int # inclusive + peak_activity: float + + @property + def duration_cols(self) -> int: + return self.end_col - self.start_col + 1 + + +def column_activity(transient_mask: np.ndarray, threshold: float = 0.5) -> np.ndarray: + """Fraction of frequency bins at or above ``threshold``, per time column.""" + return (transient_mask >= threshold).mean(axis=0) + + +def _contiguous_runs(active: np.ndarray) -> list[tuple[int, int]]: + """Inclusive (start, end) index pairs of each True run.""" + padded = np.concatenate(([False], active, [False])) + edges = np.flatnonzero(np.diff(padded.astype(np.int8))) + starts, ends = edges[::2], edges[1::2] - 1 + return list(zip(starts.tolist(), ends.tolist(), strict=True)) + + +def _fill_gaps(active: np.ndarray, max_gap: int) -> np.ndarray: + """Close False gaps of at most ``max_gap`` columns between True runs.""" + if max_gap <= 0: + return active + filled = active.copy() + runs = _contiguous_runs(active) + for (_, prev_end), (next_start, _) in zip(runs, runs[1:], strict=False): + if next_start - prev_end - 1 <= max_gap: + filled[prev_end : next_start + 1] = True + return filled + + +def extract_elm_events( + transient_mask: np.ndarray, + *, + threshold: float = 0.5, + activity_min: float = 0.1, + min_gap_cols: int = 3, + min_duration_cols: int = 1, +) -> list[ElmEvent]: + """Detect ELM events in a ``(H, W)`` transient-activity mask. + + ``threshold`` binarizes mask values; ``activity_min`` is the minimum + active-bin fraction for a column to count as part of an event; + runs separated by gaps of at most ``min_gap_cols`` columns are merged; + events shorter than ``min_duration_cols`` are dropped. + """ + activity = column_activity(transient_mask, threshold=threshold) + active = _fill_gaps(activity >= activity_min, min_gap_cols) + return [ + ElmEvent(start, end, float(activity[start : end + 1].max())) + for start, end in _contiguous_runs(active) + if end - start + 1 >= min_duration_cols + ] + + +def summarize( + events: list[ElmEvent], n_cols: int, hop: int, fs: float | None +) -> dict[str, float | int | None]: + """Per-input summary: event count, ELM frequency (needs ``fs``), duty cycle. + + ``elm_freq_hz`` is events per second of analyzed signal; ``None`` when the + sampling rate is unknown (spectrogram columns have no absolute timebase). + """ + active_cols = sum(event.duration_cols for event in events) + total_s = n_cols * hop / fs if fs else None + return { + "n_events": len(events), + "elm_freq_hz": len(events) / total_s if total_s else None, + "duty_cycle": active_cols / n_cols if n_cols else 0.0, + } + + +def _col_to_s(col: int, hop: int, fs: float | None) -> float | str: + return col * hop / fs if fs else "" + + +EVENT_FIELDS = ( + "input", + "event", + "start_col", + "end_col", + "duration_cols", + "t_start_s", + "t_end_s", + "duration_s", + "peak_activity", +) + +SUMMARY_FIELDS = ("input", "n_events", "elm_freq_hz", "duty_cycle") + + +def write_events_csv( + path: Path, + per_input: list[tuple[str, list[ElmEvent]]], + hop: int, + fs: float | None, +) -> None: + """One row per detected event; time columns blank when ``fs`` is unknown.""" + with path.open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=EVENT_FIELDS) + writer.writeheader() + for name, events in per_input: + for index, event in enumerate(events): + writer.writerow( + { + "input": name, + "event": index, + "start_col": event.start_col, + "end_col": event.end_col, + "duration_cols": event.duration_cols, + "t_start_s": _col_to_s(event.start_col, hop, fs), + "t_end_s": _col_to_s(event.end_col + 1, hop, fs), + "duration_s": _col_to_s(event.duration_cols, hop, fs), + "peak_activity": event.peak_activity, + } + ) + + +def write_summary_csv( + path: Path, per_input: list[tuple[str, dict[str, float | int | None]]] +) -> None: + """One row per input file with its :func:`summarize` result.""" + with path.open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=SUMMARY_FIELDS) + writer.writeheader() + for name, summary in per_input: + row = {"input": name, **summary} + writer.writerow({k: ("" if v is None else v) for k, v in row.items()}) diff --git a/src/tokeye/hub.py b/src/tokeye/hub.py index 1115cc1..5cb3b45 100644 --- a/src/tokeye/hub.py +++ b/src/tokeye/hub.py @@ -16,6 +16,8 @@ import torch.nn as nn from huggingface_hub import hf_hub_download +from .models.ae_tf_maskrcnn.config_ae_tf_maskrcnn import AETFMaskConfig +from .models.ae_tf_maskrcnn.model_ae_tf_maskrcnn import AETFMaskModel from .models.big_tf_unet.config_big_tf_unet import BigTFUNetConfig from .models.big_tf_unet.model_big_tf_unet import BigTFUNetModel @@ -35,17 +37,35 @@ class ModelSpec: name: str filename: str # file in the HF repo builder: Callable[[], nn.Module] + repo_id: str | None = None # None -> DEFAULT_REPO_ID (TOKEYE_HF_REPO override) +# Insertion order matters: _build_from_state_dict tries specs in order, so the +# default segmentation model must stay first — U-Net checkpoints should never +# construct the (much slower) R-CNN builder. MODEL_REGISTRY: dict[str, ModelSpec] = { "big_tf_unet": ModelSpec( "big_tf_unet", "big_tf_unet_251210.pt", lambda: BigTFUNetModel(BigTFUNetConfig()), ), + "ae_tf_maskrcnn": ModelSpec( + "ae_tf_maskrcnn", + "ae_tf_maskrcnn_251223.pt", + lambda: AETFMaskModel(AETFMaskConfig(weights=None)), + repo_id="nc1/ae_tf_maskrcnn", + ), } +def repo_for(name: str) -> str: + """Hugging Face repo a model name resolves to (for error messages).""" + spec = MODEL_REGISTRY.get(str(name)) + if spec is not None and spec.repo_id is not None: + return spec.repo_id + return DEFAULT_REPO_ID + + def resolve_device(device: str = "auto") -> str: if device == "auto": return "cuda" if torch.cuda.is_available() else "cpu" @@ -59,7 +79,7 @@ def download_model(name: str = DEFAULT_MODEL, repo_id: str | None = None) -> Pat raise ValueError( f"Unknown model {name!r}; valid names: {sorted(MODEL_REGISTRY)}" ) from exc - resolved_repo_id = repo_id or DEFAULT_REPO_ID + resolved_repo_id = repo_id or spec.repo_id or DEFAULT_REPO_ID return Path(hf_hub_download(resolved_repo_id, spec.filename)) diff --git a/src/tokeye/models/ae_tf_maskrcnn/model_ae_tf_maskrcnn.py b/src/tokeye/models/ae_tf_maskrcnn/model_ae_tf_maskrcnn.py index 57711ca..9ea0a6f 100644 --- a/src/tokeye/models/ae_tf_maskrcnn/model_ae_tf_maskrcnn.py +++ b/src/tokeye/models/ae_tf_maskrcnn/model_ae_tf_maskrcnn.py @@ -1,9 +1,6 @@ import torch import torch.nn as nn -from torchvision.models.detection import ( - MaskRCNN_ResNet50_FPN_V2_Weights, - maskrcnn_resnet50_fpn_v2, -) +from torchvision.models.detection import maskrcnn_resnet50_fpn_v2 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor from torchvision.models.detection.transform import GeneralizedRCNNTransform @@ -18,9 +15,10 @@ def __init__(self, config: AETFMaskConfig): super().__init__() self.config = config - model = maskrcnn_resnet50_fpn_v2( - weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT - ) + # config.weights=None skips the COCO download entirely; pass + # AETFMaskConfig(weights=None) when a checkpoint will overwrite + # every parameter anyway (e.g. the hub registry builder). + model = maskrcnn_resnet50_fpn_v2(weights=config.weights) in_features = model.roi_heads.box_predictor.cls_score.in_features # type: ignore[union-attr] model.roi_heads.box_predictor = FastRCNNPredictor( diff --git a/src/tokeye/modesearch/README.md b/src/tokeye/modesearch/README.md new file mode 100644 index 0000000..09cce0e --- /dev/null +++ b/src/tokeye/modesearch/README.md @@ -0,0 +1,39 @@ +# modesearch — mode database and query layer (design stage) + +Status: **descriptive text only.** No schema, storage engine, or query API has +been chosen; this document records intent so the other suite tools can grow +toward it. + +## The idea + +1. **Offline crawler / cataloguer.** A batch job walks shot archives (local + HDF5 stores, MDSplus when reachable) and runs the TokEye suite on each + shot: `big_tf_unet` masks, `modespec` toroidal mode-number fits, `elmspec` + ELM events, `alfvenspec` AE detections. Every detection becomes a mode + record in the database. +2. **Mode record.** One row per mode event — the working sketch: + shot, machine, diagnostic, time interval, frequency band, mode numbers + (n, and m when available), amplitude, confidence, detector name + version, + and a pointer back to the artifact (mask file, CSV row) that produced it. + A shared record type that all suite tools can emit is the first concrete + deliverable (see docs/ROADMAP.md, "Mode catalogue schema"). +3. **Query layer.** "Find shots with an n=2 tearing mode between 2-4 kHz + during an ELM-free period" — filters on the record fields, returning shot + lists with the matching events. Interface undecided (CLI first, probably). +4. **Consumers.** Researchers hunting for reference shots; validation studies + (mode statistics vs. campaign); and the lab's fusion-world-model shot + designer, which can learn mode occurrence statistics conditioned on + plasma parameters from the same index. + +## Relationship to shotsearch + +The sibling `shotsearch` project answers "which discharges look like this +setup?" (actuator/setup similarity). modesearch answers "which discharges +contained this MHD activity?". They meet at the shot list: a designer query +could intersect both ("shots near this setup that developed an n=1 locked +mode"). + +## Non-goals for v1 + +Real-time/inter-shot operation, cross-machine schema unification, and +automatic mode labeling beyond what the detectors already emit. diff --git a/src/tokeye/modesearch/__init__.py b/src/tokeye/modesearch/__init__.py new file mode 100644 index 0000000..7e89113 --- /dev/null +++ b/src/tokeye/modesearch/__init__.py @@ -0,0 +1,18 @@ +"""modesearch — a searchable database of detected modes (design stage). + +Nothing is implemented here yet; this package reserves the name and records +the intended design. See README.md in this directory and docs/ROADMAP.md. + +The idea: an offline crawler runs the TokEye suite (big_tf_unet masks, +modespec mode-number fits, elmspec ELM events, alfvenspec AE detections) +over shot archives and indexes every detected mode into a database — one +record per mode event: shot, machine, time interval, frequency band, mode +numbers when known, amplitude, and detector provenance. Researchers then +query it ("find shots with an n=2 tearing mode between 2-4 kHz during an +ELM-free period") instead of re-scanning raw data, and the lab's +fusion-world-model shot designer can pull mode statistics from the same +index. Complements the sibling ``shotsearch`` project, which searches +discharges by actuator/setup similarity rather than by MHD activity. +""" + +from __future__ import annotations diff --git a/src/tokeye/modespec/__init__.py b/src/tokeye/modespec/__init__.py new file mode 100644 index 0000000..0da1c0c --- /dev/null +++ b/src/tokeye/modespec/__init__.py @@ -0,0 +1,10 @@ +"""Mode-number analysis suite (the modespec family). + +- :mod:`tokeye.modespec.classic` — vendored ``pymodespec``, the Python port of + the DIII-D IDL ``modespec`` tool: Mirnov spectrograms, toroidal mode-number + fits, per-shot mode CSVs. Run via ``tokeye modespec ``. +- :mod:`tokeye.modespec.deep` — placeholder for the next-generation + single-chord mode-number identification (see its README). +""" + +from __future__ import annotations diff --git a/src/tokeye/modespec/classic/LICENSE b/src/tokeye/modespec/classic/LICENSE new file mode 100644 index 0000000..c4954f1 --- /dev/null +++ b/src/tokeye/modespec/classic/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 PlasmaControl + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/tokeye/modespec/classic/PROVENANCE.md b/src/tokeye/modespec/classic/PROVENANCE.md new file mode 100644 index 0000000..242edbe --- /dev/null +++ b/src/tokeye/modespec/classic/PROVENANCE.md @@ -0,0 +1,29 @@ +# Vendored code provenance + +- Upstream: `git@github.com:PlasmaControl/pymodespec.git` (private) +- Pinned commit: `1e0e48fc6a32f2eccbf1a88a7c21ce8480f1db8d` +- Vendored: 2026-07-06 +- License: MIT (see `LICENSE`, copied from upstream) + +## Files taken + +All eight Python modules plus the example config and license: +`modespec.py`, `generate_modes.py`, `data_utils.py`, `ece_coherence.py`, +`ece_ms_zoom.py`, `ece_sawteeth.py`, `mpi_coherence.py`, `mre_utils.py`, +`modes.yaml`, `LICENSE`. Upstream notebooks and pixi files were not vendored. + +## Local modifications + +- `generate_modes.py`: `from modespec import ...` made relative + (`from .modespec import ...`); `main()` body extracted into + `run_config(config_path) -> int` (returns failed-shot count) so the + `tokeye modespec` subcommand can call it. +- `modespec.py` (4 sites) and `data_utils.py` (1 site): the + "MDSplus not available" errors now explain where MDSplus comes from + (GA cluster / conda-forge) and that fetching needs atlas.gat.com access. +- `__init__.py` and this file are additions, not upstream files. +- Style rules are relaxed for this directory in `ruff.toml` + (vendored-code policy); correctness rules (F821 etc.) remain active. + +Re-vendoring: clone upstream at a newer commit, re-copy the files above, +re-apply the modifications in this list, and update the pinned commit. diff --git a/src/tokeye/modespec/classic/__init__.py b/src/tokeye/modespec/classic/__init__.py new file mode 100644 index 0000000..0b31b17 --- /dev/null +++ b/src/tokeye/modespec/classic/__init__.py @@ -0,0 +1,36 @@ +"""Vendored ``pymodespec`` — classic DIII-D Mirnov mode analysis. + +See PROVENANCE.md (upstream, pinned commit, local modifications) and LICENSE +(MIT) in this directory. Heavy submodules (``matplotlib.pyplot`` is imported +at module load) resolve lazily via PEP 562 so importing this package stays +cheap; ``tokeye.cli.modespec`` sets the Agg backend before touching them. +""" + +from __future__ import annotations + +from typing import Any + +_EXPORTS = { + "fetch_mirnov": "modespec", + "mode_spectrogram": "modespec", + "mode_fit_timeslice": "modespec", + "mode_svd_spectrogram": "modespec", + "plot_modespec": "modespec", + "plot_svd": "modespec", + "fetch_ece": "modespec", + "ece_mode_location": "modespec", + "load_config": "generate_modes", + "detect_modes": "generate_modes", + "run_config": "generate_modes", +} + +__all__ = list(_EXPORTS) + + +def __getattr__(name: str) -> Any: + if name in _EXPORTS: + from importlib import import_module + + module = import_module(f".{_EXPORTS[name]}", __name__) + return getattr(module, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/tokeye/modespec/classic/data_utils.py b/src/tokeye/modespec/classic/data_utils.py new file mode 100644 index 0000000..2aa09b6 --- /dev/null +++ b/src/tokeye/modespec/classic/data_utils.py @@ -0,0 +1,989 @@ +""" +data_utils.py — MDSplus fetch helpers, pkl loader, and analysis utilities. + +Split out from analysis.ipynb so every function is importable and testable +without opening a notebook. +""" + +import os, pickle +import numpy as np +from scipy import signal as scipy_signal + +try: + import MDSplus as mds + MDS_AVAILABLE = True +except ImportError: + MDS_AVAILABLE = False + +# ── Paths ────────────────────────────────────────────────────────────────────── + +ATLAS = 'atlas.gat.com' +PKL_DIR = '/fusion/projects/xpsi/transient_control/rothsteina/MRE/data' + +# TM survival model — best model: rt_model_latest_pcb (38 features, valid_loss 0.1855) +# model_23_2 (231 feat, valid_loss 0.4277): hyperparameter trial only, no companion norm files +# model_25 (42 feat, valid_loss 0.5003): use dsm_infer_cakenn.py + cake_normalizations_dict.pkl +_HERE = os.path.dirname(os.path.abspath(__file__)) +TM_MODEL_PKL = os.path.join(_HERE, 'rt_model_latest_pcb.pkl') # 38-feature PCB ensemble (best) +TM_NORM_PKL = os.path.join(_HERE, 'rt_normalizations_bms_pcb.pkl') # scalar + PCA score norms +TM_PCA_PKL = os.path.join(_HERE, 'rt_pca_components_bms_pcb.pkl') # 6-profile PCA components + +# model_25: CAKENN-format 42-feature model (valid_loss 0.5003, for comparison) +TM_CAKE_MODEL_PKL = os.path.join(_HERE, 'model_25.pkl') +TM_CAKE_NORM_PKL = os.path.join(_HERE, 'cake_normalizations_dict.pkl') +TM_CAKE_PCA_PKL = os.path.join(_HERE, 'cakenn_pca_components.pkl') +DSM_INFER_CAKE_SCRIPT = os.path.join(_HERE, 'dsm_infer_cakenn.py') + +# TAPE C model (k2c weights, what PCS actually ran during shots) +# Note: tape_rtprofile weights differ from rt_model_latest_pcb.pkl — separately trained +TAPE_C_INFER_SCRIPT = os.path.join(_HERE, 'tape_c_infer.py') +TAPE_OFFLINE_BIN = os.path.join(_HERE, 'tape_offline') + +# Radial grid used by the model (33 uniform rho_N points, same as training) +TM_PROFILE_GRID = np.linspace(0, 1, 33) + +# ── Signal definitions ───────────────────────────────────────────────────────── + +GYROTRONS = {'LEIA': 4, 'R2D2': 5, 'YODA': 8, 'NASA': 9, 'HAN': 11} + +N1RMS_CANDIDATES = ['N1RMS', 'RMS01', 'N01RMS', 'BRMSA'] + +MPI_322_SIGNALS = [ + 'MPI1A322D', 'MPI2A322D', 'MPI3A322D', 'MPI4A322D', 'MPI5A322D', + 'MPI11M322D', + 'MPI1B322D', 'MPI2B322D', 'MPI3B322D', 'MPI4B322D', 'MPI5B322D', +] + +PTDATA_OVERVIEW = { + 'ip': 'IP', # A + # echpwr is fetched via RF tree (not PTDATA — PTDATA ECHPWR units are unknown) +} + +TREE_OVERVIEW = { + 'EFIT02ER': { + 'wmhd': '\\EFIT02ER::TOP.RESULTS.AEQDSK:WMHD', # J + 'betan_efit':'\\EFIT02ER::TOP.RESULTS.AEQDSK:BETAN', + 'aminor': '\\EFIT02ER::TOP.RESULTS.AEQDSK:AMINOR', # m + 'cpasma': '\\EFIT02ER::TOP.RESULTS.GEQDSK:CPASMA', # A + 'rmaxis': '\\EFIT02ER::TOP.RESULTS.GEQDSK:RMAXIS', # m (AEQDSK:RMAXIS not in EFIT02ER) + 'betap': '\\EFIT02ER::TOP.RESULTS.AEQDSK:BETAP', + 'ipmhd': '\\EFIT02ER::TOP.RESULTS.AEQDSK:IPMHD', # A + # Extra scalars needed by TM model (post-shot equivalents of EFITRT2 inputs) + 'qmin': '\\EFIT02ER::TOP.RESULTS.AEQDSK:QMIN', + 'li': '\\EFIT02ER::TOP.RESULTS.AEQDSK:LI', + 'kappa': '\\EFIT02ER::TOP.RESULTS.AEQDSK:KAPPA', + 'tribot': '\\EFIT02ER::TOP.RESULTS.AEQDSK:TRIBOT', + 'tritop': '\\EFIT02ER::TOP.RESULTS.AEQDSK:TRITOP', + 'volume': '\\EFIT02ER::TOP.RESULTS.AEQDSK:VOLUME', # m³ + }, + 'BOLOM': {'prad': '\\BOLOM::TOP.PRAD_01.PRAD:PRAD_TOT'}, # W + 'TRANSPORT': {'h98': '\\TRANSPORT::TOP.AOT.TAU:H98Y2_OT'}, + 'NB': {'pinj': '\\NB::TOP:PINJ', # kW + 'tinj': '\\NB::TOP:TINJ'}, # N·m (NBI torque) +} + +# ── Pickle cache helpers ─────────────────────────────────────────────────────── + +def _shot_dir(shot, data_dir): + d = os.path.join(data_dir, str(shot)) + os.makedirs(d, exist_ok=True) + return d + +def _cache_path(shot, key, data_dir): + return os.path.join(_shot_dir(shot, data_dir), f'{shot}_{key}.pkl') + +def save_cache(shot, key, data, data_dir): + with open(_cache_path(shot, key, data_dir), 'wb') as f: + pickle.dump(data, f) + +def load_cache(shot, key, data_dir): + p = _cache_path(shot, key, data_dir) + if os.path.exists(p): + with open(p, 'rb') as f: + return pickle.load(f) + # Backward-compat: read old JSON cache if pkl not yet written + p_json = os.path.join(_shot_dir(shot, data_dir), f'{shot}_{key}.json') + if os.path.exists(p_json): + import json + with open(p_json) as f: + return json.load(f) + return None + +def fetch_or_load(shot, key, fetch_fn, data_dir): + """Return cached pkl if present (falls back to JSON); otherwise fetch, cache as pkl, return.""" + cached = load_cache(shot, key, data_dir) + if cached is not None: + print(f' [{shot}] loaded "{key}" from cache') + return cached + if not MDS_AVAILABLE: + raise RuntimeError( + f'MDSplus unavailable and no cache for {shot}/{key}. MDSplus ' + 'ships on the GA cluster and on conda-forge (conda install -c ' + 'conda-forge mdsplus); fetching also needs atlas.gat.com access.' + ) + print(f' [{shot}] fetching "{key}" ...') + result = fetch_fn() + save_cache(shot, key, result, data_dir) + return result + +# ── Low-level MDSplus ────────────────────────────────────────────────────────── + +def fetch_ptdata(shot, signal): + """Fetch a PTDATA time-series. Returns (data_array, time_ms_array).""" + conn = mds.Connection(ATLAS) + conn.openTree('D3D', shot) + data = np.array(conn.get(f'PTDATA("{signal}", {shot})').data()) + time = np.array(conn.get(f'DIM_OF(PTDATA("{signal}", {shot}))').data()) + conn.closeAllTrees() + if time.size > 0 and np.max(np.abs(time)) < 100: + time = time * 1e3 # s → ms + return data, time + +def fetch_tree_nodes(shot, tree, nodes_dict): + """Fetch multiple nodes from one tree. Returns {key: {data, time}}.""" + results = {} + conn = mds.Connection(ATLAS) + try: + conn.openTree(tree, shot) + for key, node in nodes_dict.items(): + try: + d = np.array(conn.get(node).data()) + t = np.array(conn.get(f'DIM_OF({node})').data()) + results[key] = {'data': d.tolist(), 'time': t.tolist()} + except Exception as e: + print(f' Warning {key} ({tree}): {e}') + conn.closeAllTrees() + except Exception as e: + print(f' Could not open {tree} for {shot}: {e}') + return results + +def ts(d, key): + """Unpack (time_array, data_array) from a {data, time} signals dict.""" + arr = d.get(key, {}) + if not arr: + return np.array([]), np.array([]) + t = np.array(arr.get('time', arr.get('time_basis', []))) + v = np.array(arr.get('data', arr.get('time_series', []))) + return t, v + +# ── Signal-set fetchers ──────────────────────────────────────────────────────── + +def fetch_ech_power(shot): + """ + Fetch total ECH power in MW. + + Primary source: RF tree, where all FPWRC nodes are in Watts. + \\RF::TOP.ECH.TOTAL:ECHPWRC — total (preferred) + \\RF::TOP.ECH.{GYRO}:EC{ABBREV}FPWRC — per gyrotron, summed as fallback + + PTDATA 'ECHPWR' is NOT used — its units are unknown (not MW, not kW). + + Returns + ------- + dict with keys 'data' (MW), 'time' (ms), 'signal' (source label). + """ + try: + conn = mds.Connection(ATLAS) + conn.openTree('RF', shot) + + # Try total node first + total_node = '\\RF::TOP.ECH.TOTAL:ECHPWRC' + try: + d = np.array(conn.get(total_node).data()) + t = np.array(conn.get(f'DIM_OF({total_node})').data()) + if d.size > 1 and d.max() > 0: + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + conn.closeAllTrees() + print(f' ECH: RF TOTAL:ECHPWRC, peak={d.max()/1e6:.3f} MW') + return {'signal': total_node, 'data': (d / 1e6).tolist(), 'time': t.tolist()} + except Exception: + pass + + # Fallback: sum per-gyrotron FPWRC nodes + time_ref = None + total_pwr = None + for gyro, _ in GYROTRONS.items(): + abbrev = gyro[:3].upper() + node = f'\\RF::TOP.ECH.{gyro}:EC{abbrev}FPWRC' + try: + d = np.array(conn.get(node).data()) + t = np.array(conn.get(f'DIM_OF({node})').data()) + if d.size < 2 or d.max() == 0: + continue + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + pwr_mw = d / 1e6 + if time_ref is None: + time_ref = t + total_pwr = pwr_mw.copy() + else: + total_pwr += np.interp(time_ref, t, pwr_mw, left=0.0, right=0.0) + except Exception: + pass + + conn.closeAllTrees() + + if total_pwr is not None: + print(f' ECH: sum of FPWRC gyrotrons, peak={total_pwr.max():.3f} MW') + return {'signal': 'RF_FPWRC_sum', 'data': total_pwr.tolist(), 'time': time_ref.tolist()} + + except Exception as e: + print(f' ECH RF tree failed: {e}') + + print(' Warning: ECH power not found in RF tree — no ECHPWR stored') + return {'signal': None, 'data': [], 'time': []} + + +def fetch_overview(shot): + results = {} + for key, sig in PTDATA_OVERVIEW.items(): + try: + d, t = fetch_ptdata(shot, sig) + results[key] = {'data': d.tolist(), 'time': t.tolist()} + except Exception as e: + print(f' Warning {key}: {e}') + results['echpwr'] = fetch_ech_power(shot) + for tree, nodes in TREE_OVERVIEW.items(): + results.update(fetch_tree_nodes(shot, tree, nodes)) + return results + +def fetch_n1rms(shot): + # Primary: MHD tree — use 5 ms smoothed version; fall back to full-rate + for node, label in [ + ('\\MHD::TOP.MIRNOV:N1RMS5', 'MHD:N1RMS5'), + ('\\MHD::TOP.MIRNOV:N1RMS', 'MHD:N1RMS'), + ]: + try: + conn = mds.Connection(ATLAS) + conn.openTree('MHD', shot) + d = np.array(conn.get(node).data()) + t = np.array(conn.get(f'DIM_OF({node})').data()) + conn.closeAllTrees() + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + if d.size > 0: + print(f' Using: {label}') + return {'signal': label, 'data': d.tolist(), 'time': t.tolist()} + except Exception: + continue + # Fallback: PTDATA candidates + for sig in N1RMS_CANDIDATES: + try: + d, t = fetch_ptdata(shot, sig) + if d.size > 0: + print(f' Using: {sig}') + return {'signal': sig, 'data': d.tolist(), 'time': t.tolist()} + except Exception: + continue + print(f' Warning: no N1RMS signal found for {shot}') + return {'signal': None, 'data': [], 'time': []} + +def fetch_n2rms(shot): + for node, label in [ + ('\\MHD::TOP.MIRNOV:N2RMS5', 'MHD:N2RMS5'), + ('\\MHD::TOP.MIRNOV:N2RMS', 'MHD:N2RMS'), + ]: + try: + conn = mds.Connection(ATLAS) + conn.openTree('MHD', shot) + d = np.array(conn.get(node).data()) + t = np.array(conn.get(f'DIM_OF({node})').data()) + conn.closeAllTrees() + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + if d.size > 0: + print(f' Using: {label}') + return {'signal': label, 'data': d.tolist(), 'time': t.tolist()} + except Exception: + continue + for sig in ['N2RMS', 'RMS02', 'N02RMS']: + try: + d, t = fetch_ptdata(shot, sig) + if d.size > 0: + print(f' Using: {sig}') + return {'signal': sig, 'data': d.tolist(), 'time': t.tolist()} + except Exception: + continue + print(f' Warning: no N2RMS signal found for {shot}') + return {'signal': None, 'data': [], 'time': []} + +def fetch_neutron_rate(shot): + """ + Fetch total neutron rate from IONS tree. + + Source: \\IONS::NEUTRONSRATE — 50 kHz, full shot, units n/s. + The IONS tree is accessible from atlas.gat.com. + + Returns + ------- + dict with keys 'signal', 'data' (n/s), 'time' (ms). + """ + node = '\\IONS::NEUTRONSRATE' + try: + conn = mds.Connection(ATLAS) + conn.openTree('IONS', shot) + d = np.asarray(conn.get(f'float({node})'), dtype=float) + t = np.asarray(conn.get(f'float(dim_of({node}))'), dtype=float) + conn.closeAllTrees() + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + if d.size > 1: + print(f' [{shot}] NEUTRONSRATE: {d.size} pts, ' + f't=[{t[0]:.0f},{t[-1]:.0f}] ms, ' + f'fs={1e3/np.median(np.diff(t)):.0f} Hz') + return {'signal': node, 'data': d.tolist(), 'time': t.tolist()} + except Exception as e: + print(f' Warning fetch_neutron_rate({shot}): {e}') + return {'signal': None, 'data': [], 'time': []} + + +def fetch_mirror_angles(shot): + """ + Fetch poloidal and toroidal mirror angles for each ECH gyrotron. + + Primary: RF tree \\RF::TOP.ECH.{GYRO}:EC{ABBREV}POLANG / AZIANG + Fallback: PTDATA GYSMPOL{N} / GYSMAZI{N} + + Keys produced: {GYRO}_pol_meas, {GYRO}_tor_meas + """ + results = {} + + # --- RF tree (primary) --- + try: + conn = mds.Connection(ATLAS) + conn.openTree('RF', shot) + for gyro in GYROTRONS: + abbrev = gyro[:3].upper() + for sfx, rfsuffix in [('pol_meas', f'EC{abbrev}POLANG'), + ('tor_meas', f'EC{abbrev}AZIANG')]: + node = f'\\RF::TOP.ECH.{gyro}:{rfsuffix}' + try: + d = np.array(conn.get(node).data()) + t = np.array(conn.get(f'DIM_OF({node})').data()) + if t.size > 0 and np.max(np.abs(t)) < 100: + t = t * 1e3 + results[f'{gyro}_{sfx}'] = {'data': d.tolist(), 'time': t.tolist()} + print(f' {gyro} {sfx}: {d.size} pts via RF tree') + except Exception as e: + print(f' Warning RF {gyro} {rfsuffix}: {e}') + conn.closeAllTrees() + except Exception as e: + print(f' RF tree unavailable: {e}') + + # --- PTDATA fallback for any missing signals --- + for gyro, idx in GYROTRONS.items(): + for sfx, sig in [('pol_meas', f'GYSMPOL{idx}'), + ('tor_meas', f'GYSMAZI{idx}')]: + key = f'{gyro}_{sfx}' + if key in results: + continue + try: + d, t = fetch_ptdata(shot, sig) + results[key] = {'data': d.tolist(), 'time': t.tolist()} + print(f' {gyro} {sfx}: {d.size} pts via PTDATA {sig}') + except Exception as e: + print(f' Warning PTDATA {sig}: {e}') + return results + +def fetch_mpi(shot, signals=None): + if signals is None: + signals = MPI_322_SIGNALS + results = {} + for sig in signals: + try: + d, t = fetch_ptdata(shot, sig) + results[sig] = {'data': d.tolist(), 'time': t.tolist()} + except Exception as e: + print(f' Warning {sig}: {e}') + return results + +# Conda Python 3.11 with torch + auton-survival available +DSM_CONDA_PY = '/fusion/projects/codes/conda/omega/envs_public/general/bin/python3' +DSM_INFER_SCRIPT = os.path.join(_HERE, 'dsm_infer_38.py') + + +def fetch_thomson_cer(shot): + """ + Fetch Te, ne, Ti, and rotation profiles from ELECTRONS/IONS trees. + + Signal priority (from common.py): + Te, ne : ELECTRONS ZIPFIT → TS blessed core → TS r00 + Ti : IONS ZIPFIT:ITEMPFIT (already keV) + rot : IONS ZIPFIT_CERA:TROTFIT → ZIPFIT:TROTFIT (km/s) + + CER tree has NOPATH on atlas — all CER-derived data fetched from IONS tree. + + Returns dict {key: {data, rho, time}} where: + data shape : (n_rho, n_t) — indexed [rho_idx, time_idx] + rho : 1-D normalised grid (0–1) + time : 1-D time axis [ms] + """ + if not MDS_AVAILABLE: + return {} + results = {} + + # ── Electron profiles (ELECTRONS tree) ────────────────────────────────── + elec_candidates = { + 'Te_keV': [ + ('\\ELECTRONS::TOP.PROFILE_FITS.ZIPFIT:ETEMPFIT', 'ELECTRONS', 1.0), # keV + ('\\ELECTRONS::TOP.TS.BLESSED.CORE.FITTED:TE', 'ELECTRONS', 1e-3), # eV + ('\\ELECTRONS::TOP.TS.REVISIONS.REVISION00.CORE:TE', 'ELECTRONS', 1e-3), # eV + ], + 'ne': [ + ('\\ELECTRONS::TOP.PROFILE_FITS.ZIPFIT:EDENSFIT', 'ELECTRONS', 1.0), + ('\\ELECTRONS::TOP.TS.BLESSED.CORE.FITTED:NE', 'ELECTRONS', 1.0), + ('\\ELECTRONS::TOP.TS.REVISIONS.REVISION00.CORE:NE', 'ELECTRONS', 1.0), + ], + } + for key, candidates in elec_candidates.items(): + for node, tree, scale in candidates: + try: + conn = mds.Connection(ATLAS) + conn.openTree(tree, shot) + d = np.array(conn.get(node).data()) + # rho and time dims — ZIPFIT: dim0=rho, dim1=time + rho = np.array(conn.get(f'DIM_OF({node}, 0)').data()) + t = np.array(conn.get(f'DIM_OF({node}, 1)').data()) + conn.closeAllTrees() + if t.size == 0 or d.size == 0: + continue + if np.max(np.abs(t)) < 100: + t = t * 1e3 # s → ms + # Ensure d is (n_rho, n_t) + if d.ndim == 2: + if d.shape[0] != len(rho) and d.shape[1] == len(rho): + d = d.T + d = d * scale + results[key] = {'data': d.tolist(), 'rho': rho.tolist(), 'time': t.tolist()} + print(f' {key}: {node} ({d.shape})') + break + except Exception: + continue + + # ── Ion profiles (IONS tree) ───────────────────────────────────────────── + ion_candidates = { + 'Ti_keV': [ + ('\\IONS::TOP.PROFILE_FITS.ZIPFIT:ITEMPFIT', 'IONS', 1.0), # keV + ('\\IONS::TOP.PROFILE_FITS.ZIPFIT_CERA:ITEMPFIT', 'IONS', 1.0), # keV + ], + 'rot_kms': [ + ('\\IONS::TOP.PROFILE_FITS.ZIPFIT_CERA:TROTFIT', 'IONS', 1.0), + ('\\IONS::TOP.PROFILE_FITS.ZIPFIT:TROTFIT', 'IONS', 1.0), + ], + } + for key, candidates in ion_candidates.items(): + for node, tree, scale in candidates: + try: + conn = mds.Connection(ATLAS) + conn.openTree(tree, shot) + d = np.array(conn.get(node).data()) + # TROTFIT/ITEMPFIT: dim_of(sig,0)=rho, dim_of(sig,1)=time (per common.py) + rho = np.array(conn.get(f'DIM_OF({node}, 0)').data()) + t = np.array(conn.get(f'DIM_OF({node}, 1)').data()) + conn.closeAllTrees() + if t.size == 0 or d.size == 0: + continue + if np.max(np.abs(t)) < 100: + t = t * 1e3 + # numpy gives (n_rho, n_t) from MDSplus; normalise to (n_rho, n_t) + if d.ndim == 2: + if d.shape[0] == len(t) and d.shape[1] == len(rho): + d = d.T # (n_t, n_rho) → (n_rho, n_t) + d = d * scale + results[key] = {'data': d.tolist(), 'rho': rho.tolist(), 'time': t.tolist()} + print(f' {key}: {node} ({d.shape})') + break + except Exception: + continue + + if not results: + print(f' Warning: no Thomson/CER profiles found for {shot}') + return results + + +def _interp_profile_to_33(sig_dict, rho33, t_ms): + """ + Interpolate a 2-D profile {data(n_rho, n_t), rho, time} to 33-pt rho33 at t_ms. + Returns zero array if data unavailable. + """ + if not sig_dict or not sig_dict.get('data'): + return np.zeros(33) + t = np.asarray(sig_dict['time']) + rho = np.asarray(sig_dict['rho']) + d = np.asarray(sig_dict['data']) + if d.ndim == 1: + return np.interp(rho33, rho, d) + # 2-D: normalise to (n_rho, n_t) + if d.shape[0] == len(t) and d.shape[1] == len(rho): + d = d.T # (n_t, n_rho) → (n_rho, n_t) + i = np.searchsorted(t, t_ms) + i = np.clip(i, 0, len(t) - 1) + return np.interp(rho33, rho, d[:, i]) + + +def run_tm_inference(shot, data_dir): + """ + Run the DSM TM survival model (rt_model_latest_pcb.pkl) for one shot. + + Uses a subprocess call to the conda python3 environment which has + torch + auton-survival installed. Loads all inputs from cached pkl files + (no live MDSplus access needed). + + Input feature vector: 38 values (14 Z-scored scalars + 6 profiles × 4 PCA components). + Preprocessing matches the rt_x_pca_bms_pcb.pkl training pipeline exactly. + + Parameters + ---------- + shot : int shot number + data_dir : str directory containing {shot}/ subdirectory with cached pkl files + + Returns + ------- + dict with 'time' (list of ms), 'data' (risk at 500 ms), 'horizons', 'risk' + """ + import subprocess, tempfile + + for path, label in [ + (TM_MODEL_PKL, 'Model pkl'), + (TM_NORM_PKL, 'Normalization pkl'), + (TM_PCA_PKL, 'PCA pkl'), + (DSM_INFER_SCRIPT,'Inference script'), + ]: + if not os.path.exists(path): + raise FileNotFoundError(f'{label} not found: {path}') + + with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False) as tmp: + out_pkl = tmp.name + + try: + result = subprocess.run( + [DSM_CONDA_PY, DSM_INFER_SCRIPT, + str(shot), data_dir, TM_MODEL_PKL, TM_NORM_PKL, TM_PCA_PKL, out_pkl], + capture_output=True, text=True, timeout=600, + ) + if result.returncode != 0: + raise RuntimeError( + f'DSM inference failed (rc={result.returncode}):\n' + f'stdout: {result.stdout[-2000:]}\n' + f'stderr: {result.stderr[-2000:]}' + ) + print(result.stdout.strip()) + with open(out_pkl, 'rb') as fh: + return pickle.load(fh) + finally: + if os.path.exists(out_pkl): + os.unlink(out_pkl) + + +def run_tm_inference_cake(shot, data_dir): + """ + Run DSM TM model_25 (42-feature CAKENN) for one shot. + + Uses dsm_infer_cakenn.py with model_25.pkl + cake_normalizations_dict.pkl + + cakenn_pca_components.pkl. Evaluates on 5 ms grid (zero-order EFIT hold). + For comparison against run_tm_inference() (rt_model_latest_pcb, valid_loss 0.1855). + model_25 valid_loss = 0.5003. + """ + import subprocess, tempfile + + for path, label in [ + (TM_CAKE_MODEL_PKL, 'CAKE model pkl'), + (TM_CAKE_NORM_PKL, 'CAKE norm pkl'), + (TM_CAKE_PCA_PKL, 'CAKE PCA pkl'), + (DSM_INFER_CAKE_SCRIPT, 'CAKE inference script'), + ]: + if not os.path.exists(path): + raise FileNotFoundError(f'{label} not found: {path}') + + with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False) as tmp: + out_pkl = tmp.name + + try: + result = subprocess.run( + [DSM_CONDA_PY, DSM_INFER_CAKE_SCRIPT, + str(shot), data_dir, + TM_CAKE_MODEL_PKL, TM_CAKE_NORM_PKL, TM_CAKE_PCA_PKL, out_pkl], + capture_output=True, text=True, timeout=600, + ) + if result.returncode != 0: + raise RuntimeError( + f'CAKE DSM inference failed (rc={result.returncode}):\n' + f'stdout: {result.stdout[-2000:]}\n' + f'stderr: {result.stderr[-2000:]}' + ) + print(result.stdout.strip()) + with open(out_pkl, 'rb') as fh: + return pickle.load(fh) + finally: + if os.path.exists(out_pkl): + os.unlink(out_pkl) + + +def run_tm_inference_c(shot, data_dir): + """ + Run TAPE C (k2c) rtProfiles model offline — same weights as PCS during the shot. + + Uses tape_c_infer.py → tape_offline binary (compiled from tape_rtprofile_*.c). + Same 38-feature BMS-PCB preprocessing as run_tm_inference(). + Outputs survival (not risk); key 'data' = survival@500ms. + + NOTE: tape_rtprofile C weights are a separately trained model from + rt_model_latest_pcb.pkl. This comparison shows PCS actual vs offline Python model. + """ + import subprocess, tempfile + + for path, label in [ + (TAPE_C_INFER_SCRIPT, 'C infer script'), + (TAPE_OFFLINE_BIN, 'tape_offline binary'), + (TM_NORM_PKL, 'Normalization pkl'), + (TM_PCA_PKL, 'PCA pkl'), + ]: + if not os.path.exists(path): + raise FileNotFoundError(f'{label} not found: {path}') + + with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False) as tmp: + out_pkl = tmp.name + + try: + result = subprocess.run( + [DSM_CONDA_PY, TAPE_C_INFER_SCRIPT, + str(shot), data_dir, TM_NORM_PKL, TM_PCA_PKL, out_pkl], + capture_output=True, text=True, timeout=600, + ) + if result.returncode != 0: + raise RuntimeError( + f'TAPE C inference failed (rc={result.returncode}):\n' + f'stdout: {result.stdout[-2000:]}\n' + f'stderr: {result.stderr[-2000:]}' + ) + print(result.stdout.strip()) + with open(out_pkl, 'rb') as fh: + return pickle.load(fh) + finally: + if os.path.exists(out_pkl): + os.unlink(out_pkl) + + +def fetch_prediction(shot, signal, data_dir): + try: + d, t = fetch_ptdata(shot, signal) + return {'data': d.tolist(), 'time': t.tolist()} + except Exception as e: + print(f' Warning prediction signal {signal!r}: {e}') + local = os.path.join(_shot_dir(shot, data_dir), f'{shot}_prediction.pkl') + if os.path.exists(local): + with open(local, 'rb') as fh: + return pickle.load(fh) + return {'data': [], 'time': []} + +# ── pkl loading ──────────────────────────────────────────────────────────────── + +class _Stub: + """Placeholder for unknown classes (e.g. xarray.DataArray) during unpickling.""" + def __init__(self, *a, **kw): pass + def __setstate__(self, s): + if isinstance(s, dict): + self.__dict__.update(s) + +class _PermissiveUnpickler(pickle.Unpickler): + def find_class(self, module, name): + try: + return super().find_class(module, name) + except (ImportError, AttributeError): + return _Stub + +def load_pkl_slice(shot, time_ms): + """ + Load one time-slice pkl from the pre-computed MRE data directory. + + Parameters + ---------- + shot : int shot number + time_ms : int time stamp in ms (must match a file name, e.g. 2500) + + Returns + ------- + dict with keys: + psiN, rhoN (129,) — normalized flux / radius grids + q_prof, p_prof (129,) — safety factor, pressure [Pa] + a_prof, Bp (129,) — minor radius [m], poloidal field [T] + R0, Bphi0, betaP — scalars (m, T, dimensionless) + eta_prof (201,) — resistivity [Ω·m] on rho_201 grid + J_BS (201,) — bootstrap current density [A/m²] + coll_i (201,192) — ion collisionality vs (rho_201, t_full) + rhostar (201,192) — ion ρ* vs (rho_201, t_full) + coll_i_dim_rho (201,) — rho grid for 201-pt arrays + coll_i_dim_t (192,) — time axis [ms] for full-shot 2-D arrays + eccd_prof (3,201) — ECCD current density per gyrotron [A/m²?] + peak_R_{0,1,2} (192,) — ECCD deposition R [m] vs time per gyrotron + peak_Z_{0,1,2} (192,) — ECCD deposition Z [m] vs time per gyrotron + peak_R_t_{0,1,2} (192,) — time axis [ms] for deposition arrays + rho_grid (129,129) — rho on (R,Z) MHD grid + R_grid, Z_grid (129,) — MHD grid [m] + ne, Te — xarray stubs (require xarray to fully load) + """ + p = os.path.join(PKL_DIR, str(shot), f'{int(time_ms)}.pkl') + with open(p, 'rb') as f: + return _PermissiveUnpickler(f).load() + +def pkl_times(shot): + """Return sorted list of available time stamps (int ms) for a shot.""" + d = os.path.join(PKL_DIR, str(shot)) + times = [] + for f in os.listdir(d): + if f.endswith('.pkl'): + stem = f[:-4] + if stem.isdigit(): # skip cache files like {shot}_{key}.pkl + times.append(int(stem)) + return sorted(times) + +def load_all_pkl(shot, verbose=False): + """ + Load all pkl time-slices for one shot. + + Returns + ------- + dict {time_ms (int): data_dict} + """ + times = pkl_times(shot) + slices = {} + for t in times: + slices[t] = load_pkl_slice(shot, t) + if verbose: + print(f' loaded t={t} ms') + return slices + +def pkl_scalar_series(slices, key): + """Extract a scalar field (R0, betaP, …) across all time slices.""" + times = sorted(slices) + vals = [float(slices[t][key]) for t in times] + return np.array(times, dtype=float), np.array(vals) + +def pkl_profile_at_rho(slices, key, rho_target, rho_key='rhoN', grid_201=False): + """ + Interpolate a 1-D profile field to rho_target at each time slice. + + Parameters + ---------- + key : field name (e.g. 'q_prof', 'eta_prof', 'J_BS') + rho_target : float normalized rho at which to evaluate + rho_key : 'rhoN' for 129-pt arrays; 'coll_i_dim_rho' for 201-pt arrays + grid_201 : if True use coll_i_dim_rho (201 pts) instead of rhoN + + Returns + ------- + times_ms (N,), values (N,) + """ + times = sorted(slices) + vals = [] + for t in times: + d = slices[t] + v = np.asarray(d[key]) + rho = np.asarray(d['coll_i_dim_rho' if grid_201 else rho_key]) + vals.append(float(np.interp(rho_target, rho, v))) + return np.array(times, dtype=float), np.array(vals) + +# ── MRE equilibrium helpers ──────────────────────────────────────────────────── + +def find_q_surface_pkl(d, q_target=2.0, rho_min=0.25): + """ + Find minor radius r_s where q = q_target from a pkl slice. + + Returns (rho_s, r_s) or (None, None) if q_target not crossed. + rho_s is normalized (0–1), r_s is in metres. + + Takes the outermost crossing with rho_s > rho_min to skip reversed-shear + inner crossings and axis-grazing spikes. + """ + rhoN = np.asarray(d['rhoN']) + q_prof = np.asarray(d['q_prof']) + a_prof = np.asarray(d['a_prof']) + crossings = np.where(np.diff(np.sign(q_prof - q_target)))[0] + if len(crossings) == 0: + return None, None + # resolve each crossing to rho and keep only those outside rho_min + valid_i = [] + for idx in crossings: + frac = (q_target - q_prof[idx]) / (q_prof[idx+1] - q_prof[idx]) + rho_c = rhoN[idx] + frac * (rhoN[idx+1] - rhoN[idx]) + if rho_c >= rho_min: + valid_i.append(idx) + if not valid_i: + return None, None + i = valid_i[-1] # outermost qualifying crossing + frac = (q_target - q_prof[i]) / (q_prof[i+1] - q_prof[i]) + rho_s = rhoN[i] + frac * (rhoN[i+1] - rhoN[i]) + r_s = float(np.interp(rho_s, rhoN, a_prof)) + return float(rho_s), r_s + +def mre_quantities_from_pkl(d, q_target=2.0): + """ + Compute MRE-relevant scalar quantities from one pkl time slice. + + Returns dict with: + rho_s, r_s — q=q_target surface location + Bp_s — poloidal field at r_s [T] + eta_s — resistivity at r_s [Ω·m] + tau_R — resistive time μ₀r_s²/(1.22η) [s] + J_BS_s — bootstrap current density at r_s [A/m²] + betaP — scalar poloidal beta + Lp, Lq — pressure and q scale lengths [m] (None if not computable) + eccd_at_s (3,) — ECCD current density per gyrotron at r_s [A/m²] + Returns None if q=2 surface not found. + """ + MU0 = 4 * np.pi * 1e-7 + + rhoN = np.asarray(d['rhoN']) + q_prof = np.asarray(d['q_prof']) + p_prof = np.asarray(d['p_prof']) + a_prof = np.asarray(d['a_prof']) + Bp = np.asarray(d['Bp']) + + # a_prof is only populated for the inner fraction of the grid; extrapolate + # linearly through the valid (non-zero) points to get r(rho) everywhere. + valid = a_prof > 1e-6 + if valid.sum() >= 2: + slope, intercept = np.polyfit(rhoN[valid], a_prof[valid], 1) + r_phys = np.clip(rhoN * slope + intercept, 0.0, None) + else: + r_phys = a_prof.copy() + + rho_s, _ = find_q_surface_pkl(d, q_target) + if rho_s is None: + return None + r_s = float(np.interp(rho_s, rhoN, r_phys)) + + rho201 = np.asarray(d['coll_i_dim_rho']) + eta = np.asarray(d['eta_prof']) + J_BS = np.asarray(d['J_BS']) + + Bp_s = float(np.interp(rho_s, rhoN, Bp)) + eta_s = float(np.interp(rho_s, rho201, eta)) + tau_R = MU0 * r_s**2 / (1.22 * eta_s) if eta_s > 0 else None + J_BS_s = float(np.interp(rho_s, rho201, J_BS)) + + # Scale lengths at r_s. Lp = -p/(dp/dr) > 0 (p decreasing outward). + # Lq = +q/(dq/dr) > 0 (q increasing outward) — note POSITIVE sign. + def _L(prof, sign=-1): + ddr = np.gradient(prof, r_phys) + f_s = float(np.interp(rho_s, rhoN, prof)) + df_s = float(np.interp(rho_s, rhoN, ddr)) + if df_s == 0 or f_s == 0: + return None + L = abs(sign * f_s / df_s) + return L if L > 1e-10 else None + Lp = _L(p_prof, sign=-1) # pressure decreases → df_s < 0 → -f/df > 0 + Lq = _L(q_prof, sign=+1) # q increases → df_s > 0 → +f/df > 0 + + # ECCD current density at q=2 surface (3 gyrotrons × 201-pt rho grid) + eccd = np.asarray(d['eccd_prof']) # (3, 201) + eccd_at_s = np.array([float(np.interp(rho_s, rho201, eccd[i])) for i in range(3)]) + + return { + 'rho_s': rho_s, + 'r_s': r_s, + 'Bp_s': Bp_s, + 'eta_s': eta_s, + 'tau_R': tau_R, + 'J_BS_s': J_BS_s, + 'betaP': float(d['betaP']), + 'Lp': Lp, + 'Lq': Lq, + 'eccd_at_s': eccd_at_s, + } + +def mre_timeseries(slices, q_target=2.0): + """ + Compute mre_quantities_from_pkl for every time slice. + + Returns + ------- + times_ms : ndarray (N,) + mre_ts : list of dicts (or None where q surface not found) + """ + times = sorted(slices) + results = [] + for t in times: + results.append(mre_quantities_from_pkl(slices[t], q_target)) + return np.array(times, dtype=float), results + +# ── Analysis utilities ───────────────────────────────────────────────────────── + +def check_tm_events(data, time, thresh_G, min_dur_ms, smooth_ms=100.0): + """Find segments where the rolling-maximum envelope of |data| exceeds thresh_G + for >= min_dur_ms. smooth_ms sets the rolling-max window to handle the + amplitude modulation of a rotating mode (N1RMS at 5 ms resolution oscillates + at the rotation frequency, so raw contiguous threshold checks fail).""" + data, time = np.asarray(data, float), np.asarray(time, float) + amp = np.abs(data) + + # Causal rolling maximum over smooth_ms look-back window. + # Envelope[i] = max(amp[i-w:i+1]) so the envelope rises as soon as the + # signal crosses threshold, avoiding a non-causal shift of the onset time. + dt_ms = float(np.median(np.diff(time))) if len(time) > 1 else 5.0 + w = max(1, int(round(smooth_ms / dt_ms))) + n = len(amp) + envelope = np.array([amp[max(0, i - w):i + 1].max() for i in range(n)]) + + above = envelope > thresh_G + events, in_event, t_start, i_start = [], False, None, None + for i in range(len(time)): + if above[i] and not in_event: + in_event, t_start, i_start = True, time[i], i + elif not above[i] and in_event: + dur = time[i-1] - t_start + if dur >= min_dur_ms: + events.append({'t_start': float(t_start), 't_end': float(time[i-1]), + 'duration_ms': float(dur), + 'peak_G': float(amp[i_start:i].max())}) + in_event = False + if in_event: + dur = time[-1] - t_start + if dur >= min_dur_ms: + events.append({'t_start': float(t_start), 't_end': float(time[-1]), + 'duration_ms': float(dur), + 'peak_G': float(amp[i_start:].max())}) + return events + +def get_tm_onset(n1rms_dict, t_win, thresh_G, min_dur_ms): + data = np.array(n1rms_dict.get('data', [])) + time = np.array(n1rms_dict.get('time', [])) + if data.size == 0: + return None + mask = (time >= t_win[0]) & (time <= t_win[1]) + evts = check_tm_events(data[mask], time[mask], thresh_G, min_dur_ms) + return evts[0]['t_start'] if evts else None + +def band_rms(data, time_ms, f_lo=5e3, f_hi=25e3): + """Band-pass filter (f_lo–f_hi Hz) and return instantaneous amplitude.""" + dt = np.median(np.diff(time_ms)) * 1e-3 # ms → s + fs = 1.0 / dt + nyq = 0.5 * fs + lo, hi = f_lo / nyq, min(f_hi / nyq, 0.99) + if lo >= hi or lo <= 0: + return np.abs(data) + b, a = scipy_signal.butter(4, [lo, hi], btype='band') + return np.abs(scipy_signal.filtfilt(b, a, data)) + +def compute_br_tilde(mpi_dict, t_win): + """ + RMS envelope of band-filtered MPI signals as a proxy for |B̃_r| [Gauss]. + Returns (time_ms, amplitude_G). + """ + all_env, t_ref = [], None + for arr in mpi_dict.values(): + t = np.array(arr.get('time', [])) + v = np.array(arr.get('data', [])) + if t.size == 0: + continue + mask = (t >= t_win[0]) & (t <= t_win[1]) + if mask.sum() < 10: + continue + all_env.append(band_rms(v[mask], t[mask])) + if t_ref is None: + t_ref = t[mask] + if not all_env or t_ref is None: + return np.array([]), np.array([]) + n = min(len(r) for r in all_env) + return t_ref[:n], np.sqrt(np.mean(np.vstack([r[:n] for r in all_env])**2, axis=0)) diff --git a/src/tokeye/modespec/classic/ece_coherence.py b/src/tokeye/modespec/classic/ece_coherence.py new file mode 100644 index 0000000..44d5285 --- /dev/null +++ b/src/tokeye/modespec/classic/ece_coherence.py @@ -0,0 +1,282 @@ +""" +ECE cross-channel coherence and phase — quiet vs active EPM phases. + +Tests whether Te fluctuations present during 'quiet' phases are: + (a) spatially coherent across radii → real sub-threshold mode + (b) phase-locked between channels → confirms EPM spatial structure + (c) change in phase relationship between quiet and active phases + +ECE is at 5 kHz (0.2ms). EPM at 8-20 kHz aliases to: + 8 kHz → 2.0 kHz, 9 kHz → 1.0 kHz, 10 kHz → DC, + 11 kHz → 1.0 kHz, 12 kHz → 2.0 kHz, 20 kHz → DC + +Run: /fusion/projects/codes/conda/omega/envs_public/general/bin/python3 ece_coherence.py +""" +import sys, os, json +sys.path.insert(0, '/home/yasodak/NTM_premptive_control') + +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from scipy.signal import coherence, welch, csd + +SHOTS = [199606, 199607] +LONG_EVENT = {199606: (2746, 4056), 199607: (3496, 4801)} + +# Channels spanning core to edge; ECE22~q=2, ECE38~q_min, ECE46~core +CHANNELS = [22, 26, 30, 34, 38, 42, 46] +LABELS = ['~q=2','R~1.90','R~1.85','R~1.80','~q_min','R~1.73','core'] +REF_CH = 38 # reference channel (q_min region — most active) + +FS = 5000.0 # Hz (0.2ms cadence) +DT = 1.0/FS # s + +def load_ece(shot, ch): + cp = f'/home/yasodak/exp/{shot}/ptdata_ECE{ch:02d}_{shot}.json' + with open(cp) as f: d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: t = t*1e3 + return t, y + +def load_n1rms(shot): + cp = f'/home/yasodak/exp/{shot}/mds_mhd_MHD__TOP_MIRNOV_N1RMS5.json' + with open(cp) as f: d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: t = t*1e3 + return t, y + +def get_window(t, y, tmin, tmax): + mk = (t >= tmin) & (t <= tmax) & np.isfinite(y) + return y[mk] + +# Segment definitions: quiet = pre-long-event, active = long event +SEGS = { + # 199606: inter-burst quiet gaps within the burst period are < 40ms (< 200 pts at 5 kHz), + # too short for frequency-resolved coherence. Use pre-onset flatop as quiet reference. + # 199607: quiet gap at ~3400ms is 115ms = 575 pts, insufficient for NPERSEG=1000. + # Use pre-onset (2200–3370ms) as the ECE quiet reference. + # The MPI analysis (200 kHz) uses the actual inter-burst gaps at 3400ms and 4800ms. + 199606: { + 'quiet': (2000, 2740), # pre-onset quiet (best available for ECE 5kHz) + 'active': (3200, 3800), # burst period — N1RMS5 peaks 13–20 G, 600ms = 3000 pts + }, + 199607: { + 'quiet': (2200, 3360), # pre-onset (inter-burst gap at 3400ms too short for ECE) + 'active': (3800, 4700), # main burst period (N1RMS5 12–28 G) + }, +} + +# nperseg for Welch/coherence: 0.2s → 1000 pts (shorter segments = more averages) +NPERSEG = 1000 + +# Aliased EPM frequencies (EPM at 8–20 kHz aliases into ECE at 5 kHz) +EPM_ALIASES = { + '1kHz (9/11 kHz)': (800, 1200), + '2kHz (8/12/18 kHz)': (1700, 2200), +} +EPM_ALIAS_CENTERS = [1000, 2000] # Hz — for vertical markers + +for shot in SHOTS: + print(f'\n=== Shot {shot} ===') + ev = LONG_EVENT[shot] + segs = SEGS[shot] + + # Load all ECE channels + ece = {} + t_ref = None + for ch in CHANNELS: + try: + t, y = load_ece(shot, ch) + ece[ch] = (t, y) + if t_ref is None: + t_ref = t + except Exception as e: + print(f' ECE{ch}: not available') + + t_n1, n1 = load_n1rms(shot) + + # ── Figure: 3 columns = quiet / active / difference; rows = each channel pair ── + fig, axes = plt.subplots(len(CHANNELS), 3, figsize=(15, 2.5*len(CHANNELS)), + gridspec_kw={'wspace': 0.3, 'hspace': 0.1}) + fig.suptitle(f'Shot {shot}: ECE cross-coherence with ECE{REF_CH} — quiet vs active', + fontsize=11) + + col_titles = [f"Quiet {segs['quiet']}ms", f"Active {segs['active']}ms", + 'Coherence difference (active − quiet)'] + for c, ttl in enumerate(col_titles): + axes[0, c].set_title(ttl, fontsize=9) + + for row, (ch, lbl) in enumerate(zip(CHANNELS, LABELS)): + ax_q = axes[row, 0] + ax_a = axes[row, 1] + ax_d = axes[row, 2] + + if ch not in ece or REF_CH not in ece: + continue + + t_ch, y_ch = ece[ch] + t_ref_ch, y_ref = ece[REF_CH] + + for ax, (tmin, tmax), color, phase_label in [ + (ax_q, segs['quiet'], 'tab:green', 'quiet'), + (ax_a, segs['active'], 'tab:red', 'active'), + ]: + # extract segments on same time grid (both ECE channels same rate) + mk = (t_ch >= tmin) & (t_ch <= tmax) & np.isfinite(y_ch) + mk_r = (t_ref_ch >= tmin) & (t_ref_ch <= tmax) & np.isfinite(y_ref) + n_pts = min(mk.sum(), mk_r.sum()) + if n_pts < NPERSEG * 2: + ax.text(0.5, 0.5, 'insufficient data', transform=ax.transAxes, + ha='center', fontsize=7) + continue + y1 = y_ch[mk][:n_pts] + y2 = y_ref[mk_r][:n_pts] + + # Cross-coherence + f_c, coh = coherence(y1, y2, fs=FS, nperseg=NPERSEG) + # Cross-spectral phase + f_s, Pxy = csd(y1, y2, fs=FS, nperseg=NPERSEG) + phase_xy = np.angle(Pxy, deg=True) + + ax.plot(f_c, coh, color=color, lw=0.8) + ax.set_ylim(0, 1) + ax.set_xlim(0, 2500) + + # Mark aliased EPM frequencies + for f_epm in [8000, 9000, 10000, 11000, 12000, 15000, 18000, 20000]: + f_alias = abs((f_epm % int(FS)) - (FS if (f_epm % int(FS)) > FS/2 else 0)) + if f_alias > 0: + ax.axvline(f_alias, color='gray', lw=0.4, ls=':', alpha=0.5) + + if row == len(CHANNELS)-1: + ax.set_xlabel('f (Hz)', fontsize=6) + ax.tick_params(labelsize=5) + if ch == REF_CH: + ax.set_ylim(0, 1) + + if row == 0 or True: + axes[row, 0].set_ylabel(f'ECE{ch} {lbl}\nvs ECE{REF_CH}', fontsize=6) + + # Difference panel: coherence active - quiet + try: + tmin_q, tmax_q = segs['quiet'] + tmin_a, tmax_a = segs['active'] + mk_q = (t_ch >= tmin_q) & (t_ch <= tmax_q) & np.isfinite(y_ch) + mk_a = (t_ch >= tmin_a) & (t_ch <= tmax_a) & np.isfinite(y_ch) + mk_rq = (t_ref_ch >= tmin_q) & (t_ref_ch <= tmax_q) & np.isfinite(y_ref) + mk_ra = (t_ref_ch >= tmin_a) & (t_ref_ch <= tmax_a) & np.isfinite(y_ref) + n_q = min(mk_q.sum(), mk_rq.sum()) + n_a = min(mk_a.sum(), mk_ra.sum()) + if n_q >= NPERSEG*2 and n_a >= NPERSEG*2: + _, coh_q = coherence(y_ch[mk_q][:n_q], y_ref[mk_rq][:n_q], fs=FS, nperseg=NPERSEG) + f_a, coh_a = coherence(y_ch[mk_a][:n_a], y_ref[mk_ra][:n_a], fs=FS, nperseg=NPERSEG) + diff = coh_a - coh_q + ax_d.plot(f_a, diff, color='k', lw=0.8) + ax_d.axhline(0, color='gray', lw=0.5) + ax_d.fill_between(f_a, diff, 0, + where=diff > 0, color='tab:red', alpha=0.4) + ax_d.fill_between(f_a, diff, 0, + where=diff < 0, color='tab:green', alpha=0.4) + ax_d.set_ylim(-1, 1) + ax_d.set_xlim(0, 2500) + # Find frequency of max coherence increase + if len(f_a) > 0: + peak = f_a[np.argmax(diff)] + ax_d.axvline(peak, color='tab:red', lw=0.8, ls='--') + print(f' ECE{ch} vs ECE{REF_CH}: max coherence gain at {peak:.0f} Hz ' + f'(Δcoh={np.max(diff):.2f})') + if row == len(CHANNELS)-1: + ax_d.set_xlabel('f (Hz)', fontsize=6) + ax_d.tick_params(labelsize=5) + except Exception as e: + ax_d.text(0.5, 0.5, str(e)[:30], transform=ax_d.transAxes, fontsize=6) + + out = f'figures/ece_coherence_{shot}.png' + fig.savefig(out, dpi=120, bbox_inches='tight') + print(f' Saved {out}') + plt.close(fig) + + # ── Summary figure: coherence at alias bands vs channel ────────────────── + fig2, axes2 = plt.subplots(2, 2, figsize=(11, 7), + gridspec_kw={'hspace': 0.35, 'wspace': 0.3}) + fig2.suptitle(f'Shot {shot}: ECE coherence at EPM alias frequencies vs channel\n' + f'(ref = ECE{REF_CH}, quiet={segs["quiet"]}ms, active={segs["active"]}ms)', + fontsize=10) + + # For each alias band collect coherence per channel + alias_bands = [(800, 1200, '~1 kHz\n(9/11 kHz EPM alias)'), + (1700, 2300, '~2 kHz\n(8/12/18 kHz EPM alias)')] + + # Also compute full-band max coherence per channel + coh_summary = {ch: {'quiet': {}, 'active': {}} for ch in CHANNELS} + + for ch in CHANNELS: + if ch not in ece or REF_CH not in ece: + continue + t_ch, y_ch = ece[ch] + t_ref_ch, y_ref = ece[REF_CH] + for phase_key, (tmin, tmax) in [('quiet', segs['quiet']), ('active', segs['active'])]: + mk = (t_ch >= tmin) & (t_ch <= tmax) & np.isfinite(y_ch) + mk_r = (t_ref_ch >= tmin) & (t_ref_ch <= tmax) & np.isfinite(y_ref) + n_pts = min(mk.sum(), mk_r.sum()) + if n_pts < NPERSEG * 2: + continue + f_c, coh = coherence(y_ch[mk][:n_pts], y_ref[mk_r][:n_pts], fs=FS, nperseg=NPERSEG) + for (flo, fhi, _) in alias_bands: + band = (f_c >= flo) & (f_c <= fhi) + coh_summary[ch][phase_key][(flo, fhi)] = np.mean(coh[band]) if band.sum() > 0 else np.nan + # Broadband 50–2000 Hz + band_bb = (f_c >= 50) & (f_c <= 2000) + coh_summary[ch][phase_key]['broad'] = np.mean(coh[band_bb]) if band_bb.sum() > 0 else np.nan + + x_pos = np.arange(len(CHANNELS)) + bar_w = 0.35 + + for ai, (flo, fhi, band_lbl) in enumerate(alias_bands): + ax_b = axes2[ai, 0] + ax_d2 = axes2[ai, 1] + + coh_q_arr = [coh_summary[ch]['quiet'].get((flo, fhi), np.nan) for ch in CHANNELS] + coh_a_arr = [coh_summary[ch]['active'].get((flo, fhi), np.nan) for ch in CHANNELS] + + ax_b.bar(x_pos - bar_w/2, coh_q_arr, bar_w, color='tab:green', alpha=0.8, + label=f'Quiet {segs["quiet"]}ms') + ax_b.bar(x_pos + bar_w/2, coh_a_arr, bar_w, color='tab:red', alpha=0.8, + label=f'Active {segs["active"]}ms') + ax_b.set_xticks(x_pos) + ax_b.set_xticklabels([f'ECE{c}\n{l}' for c, l in zip(CHANNELS, LABELS)], + fontsize=6, rotation=30) + ax_b.set_ylim(0, 1) + ax_b.set_ylabel('Mean coherence', fontsize=8) + ax_b.set_title(band_lbl, fontsize=8) + ax_b.legend(fontsize=7) + ax_b.axhline(2/NPERSEG * np.log(20), color='k', lw=0.8, ls='--', label='95% sig.') + ax_b.tick_params(labelsize=6) + + diff_arr = np.array(coh_a_arr) - np.array(coh_q_arr) + ax_d2.bar(x_pos, diff_arr, + color=['tab:red' if d > 0 else 'tab:green' for d in diff_arr]) + ax_d2.axhline(0, color='k', lw=0.8) + ax_d2.set_xticks(x_pos) + ax_d2.set_xticklabels([f'ECE{c}' for c in CHANNELS], fontsize=7, rotation=30) + ax_d2.set_ylim(-1, 1) + ax_d2.set_ylabel('Δcoh (active − quiet)', fontsize=8) + ax_d2.set_title(f'Coherence gain {band_lbl}', fontsize=8) + ax_d2.tick_params(labelsize=6) + + # Print summary + print(f' [{band_lbl.split(chr(10))[0]}] coherence per channel:') + for ch, cq, ca in zip(CHANNELS, coh_q_arr, coh_a_arr): + dc = ca - cq if not (np.isnan(ca) or np.isnan(cq)) else np.nan + sq = f'{cq:.3f}' if not np.isnan(cq) else '---' + sa = f'{ca:.3f}' if not np.isnan(ca) else '---' + sd = f'{dc:.3f}' if not np.isnan(dc) else '---' + print(f' ECE{ch}: quiet={sq}, active={sa}, Δ={sd}') + + out2 = f'figures/ece_coherence_summary_{shot}.png' + fig2.savefig(out2, dpi=130, bbox_inches='tight') + print(f' Saved {out2}') + plt.close(fig2) + +print('\nDone.') diff --git a/src/tokeye/modespec/classic/ece_ms_zoom.py b/src/tokeye/modespec/classic/ece_ms_zoom.py new file mode 100644 index 0000000..405268a --- /dev/null +++ b/src/tokeye/modespec/classic/ece_ms_zoom.py @@ -0,0 +1,172 @@ +""" +Millisecond-resolution ECE zoom to check sawtooth sign inversion and +crash character (quiet vs active EPM phases). + +Run: /fusion/projects/codes/conda/omega/envs_public/general/bin/python3 ece_ms_zoom.py +""" +import sys, os, json +sys.path.insert(0, '/home/yasodak/NTM_premptive_control') + +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +SHOTS = [199606, 199607] +LONG_EVENT = {199606: (2746, 4056), 199607: (3496, 4801)} + +# ECE channels to plot (inner→outer): core is ECE46/42, q=2 ~ ECE22 +CHANNELS_MS = [46, 42, 38, 34, 30, 26, 22] +LABELS_MS = ['core','R~1.73','q_min','R~1.80','R~1.85','R~1.90','~q=2'] + +# Time windows: (tmin, tmax, label) — 50–150 ms wide for crash visibility +WINDOWS = { + 199606: [ + (2600, 2760, 'quiet before onset'), + (2750, 2900, 'at mode onset (2746ms)'), + (3200, 3400, 'active window'), + (4100, 4300, 'after long event'), + ], + 199607: [ + (3200, 3400, 'quiet before onset'), + (3450, 3600, 'at mode onset (3496ms)'), + (3700, 3900, 'active window'), + (4850, 5050, 'after long event'), + ], +} + +def load_ece(shot, ch): + cp = f"/home/yasodak/exp/{shot}/ptdata_ECE{ch:02d}_{shot}.json" + with open(cp) as f: + d = json.load(f) + return np.array(d['t']), np.array(d['y']) + +def load_n1rms(shot): + cp = f"/home/yasodak/exp/{shot}/mds_mhd_MHD__TOP_MIRNOV_N1RMS5.json" + with open(cp) as f: + d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: + t = t * 1e3 + return t, y + + +for shot in SHOTS: + print(f"\n=== Shot {shot} ===") + ev = LONG_EVENT[shot] + + ece = {} + for ch in CHANNELS_MS: + try: + t, y = load_ece(shot, ch) + ece[ch] = (t, y) + except: + print(f" ECE{ch} not cached") + + t_n1, n1 = load_n1rms(shot) + wins = WINDOWS[shot] + + fig, axes = plt.subplots(len(CHANNELS_MS)+1, len(wins), + figsize=(4.5*len(wins), 2*(len(CHANNELS_MS)+1)), + gridspec_kw={'hspace': 0.04, 'wspace': 0.08}) + fig.suptitle(f'Shot {shot}: ECE ms-zoom (core→edge) — sawtooth check', fontsize=11) + + cmap = plt.cm.plasma + for col, (tmin, tmax, title) in enumerate(wins): + in_ev = ev[0] <= (tmin+tmax)/2 <= ev[1] + + # N1RMS row + ax = axes[0, col] + mk = (t_n1 >= tmin) & (t_n1 <= tmax) + ax.plot(t_n1[mk], n1[mk], 'k', lw=0.8) + ax.axhline(12, color='red', lw=0.7, ls='--') + if in_ev: + ax.set_facecolor('#fff0f0') + ax.set_xlim(tmin, tmax) + ax.set_title(title, fontsize=7, pad=2, + color='darkred' if in_ev else 'black') + ax.tick_params(labelbottom=False, labelsize=6) + if col == 0: + ax.set_ylabel('N1RMS\n(G)', fontsize=6) + + for row, (ch, lbl) in enumerate(zip(CHANNELS_MS, LABELS_MS)): + ax2 = axes[row+1, col] + if ch in ece: + t_e, y_e = ece[ch] + mk = (t_e >= tmin) & (t_e <= tmax) + color = cmap(row / len(CHANNELS_MS)) + ax2.plot(t_e[mk], y_e[mk], lw=0.7, color=color) + if in_ev: + ax2.set_facecolor('#fff0f0') + ax2.set_xlim(tmin, tmax) + ax2.tick_params(labelsize=5) + if col == 0: + ax2.set_ylabel(f'ECE{ch}\n{lbl}', fontsize=6) + if row < len(CHANNELS_MS)-1: + ax2.tick_params(labelbottom=False) + else: + ax2.set_xlabel('t (ms)', fontsize=6) + + out = f'figures/ece_ms_zoom_{shot}.png' + fig.savefig(out, dpi=140, bbox_inches='tight') + print(f" Saved {out}") + plt.close(fig) + + # ── Also: single-crash zoom — find a crash in the core channel and zoom 10ms ── + # Detect crashes as rapid drops in ECE46 > 2*noise_level in <2ms + if 46 in ece: + t_c, y_c = ece[46] + mask = (t_c >= 2200) & (t_c <= 5400) + tc, yc = t_c[mask], y_c[mask] + dy = np.diff(yc) + noise = np.std(dy) + # Crashes: large negative dy in one step + crash_idx = np.where(dy < -5*noise)[0] + crash_times = tc[crash_idx] + print(f" ECE46 crashes detected: {len(crash_times)}") + if len(crash_times) > 0: + print(f" First 10: {crash_times[:10].astype(int).tolist()}") + + # Pick one crash in quiet phase and one in active phase + quiet_crashes = crash_times[(crash_times < ev[0]) | (crash_times > ev[1])] + active_crashes = crash_times[(crash_times >= ev[0]) & (crash_times <= ev[1])] + + sample_crashes = [] + if len(quiet_crashes) > 2: + sample_crashes.append((quiet_crashes[len(quiet_crashes)//2], 'quiet phase')) + if len(active_crashes) > 2: + sample_crashes.append((active_crashes[len(active_crashes)//2], 'active phase')) + + if sample_crashes: + fig3, axes3 = plt.subplots(len(CHANNELS_MS), len(sample_crashes), + figsize=(5*len(sample_crashes), 2*len(CHANNELS_MS)), + gridspec_kw={'hspace': 0.05, 'wspace': 0.1}, + sharex='col') + fig3.suptitle(f'Shot {shot}: single crash zoom (±15 ms)', fontsize=11) + + for col, (tc_crash, phase_lbl) in enumerate(sample_crashes): + for row, (ch, lbl) in enumerate(zip(CHANNELS_MS, LABELS_MS)): + ax3 = axes3[row, col] if len(sample_crashes) > 1 else axes3[row] + if ch in ece: + t_e, y_e = ece[ch] + mk = (t_e >= tc_crash-15) & (t_e <= tc_crash+15) + ax3.plot(t_e[mk], y_e[mk], lw=0.8, + color=cmap(row / len(CHANNELS_MS))) + ax3.axvline(tc_crash, color='gray', lw=0.8, ls='--') + ax3.set_xlim(tc_crash-15, tc_crash+15) + ax3.tick_params(labelsize=6) + if row == 0: + ax3.set_title(f'{phase_lbl}\ncrash at {tc_crash:.1f} ms', fontsize=8) + if col == 0: + ax3.set_ylabel(f'ECE{ch} {lbl}', fontsize=6) + if row < len(CHANNELS_MS)-1: + ax3.tick_params(labelbottom=False) + else: + ax3.set_xlabel('t (ms)', fontsize=6) + + out3 = f'figures/ece_crash_zoom_{shot}.png' + fig3.savefig(out3, dpi=150, bbox_inches='tight') + print(f" Saved {out3}") + plt.close(fig3) + +print("\nDone.") diff --git a/src/tokeye/modespec/classic/ece_sawteeth.py b/src/tokeye/modespec/classic/ece_sawteeth.py new file mode 100644 index 0000000..29a3e35 --- /dev/null +++ b/src/tokeye/modespec/classic/ece_sawteeth.py @@ -0,0 +1,153 @@ +""" +ECE Te time traces for sawtooth/quasi-sawtooth check. +Channels ECE10-48 via PTDATA, 0.2ms cadence. + +B0 ~ 2.0T at R0=1.7m → f_2nd_harmonic = 2*28*B GHz + - Core (R~1.70m, B~2.0T): f ~ 112 GHz → channel ~39 + - q_min (R~1.75m, B~1.94T): f ~ 109 GHz → channel ~36 + - q=2 (R~2.00m, B~1.70T): f ~ 95 GHz → channel ~22 + - Edge (R~2.25m, B~1.51T): f ~ 85 GHz → channel ~12 + +Run: /fusion/projects/codes/conda/omega/envs_public/general/bin/python3 ece_sawteeth.py +""" +import sys, os, json, subprocess +sys.path.insert(0, '/home/yasodak/NTM_premptive_control') + +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +SHOTS = [199606, 199607] +# Channels chosen to span core → edge +# (higher channel number ~ higher freq ~ smaller R ~ hotter core) +CHANNELS = [12, 18, 22, 26, 30, 34, 38, 42, 46] +LABELS = ['~edge','R~2.1m','~q=2','R~1.9m','R~1.85m', + 'R~1.80m','~q_min','R~1.73m','core'] +LONG_EVENT = {199606: (2746, 4056), 199607: (3496, 4801)} + +# Zoom windows: (wide_tmin, wide_tmax, tight_tmin, tight_tmax) +ZOOMS = { + 199606: [(2600, 2900, 2700, 2800)], # around onset at 2746ms + 199607: [(3300, 3650, 3450, 3560)], # around onset at 3496ms +} + +def fetch_ece(shot, ch): + key = f"ptdata_ECE{ch:02d}_{shot}" + cp = f"/home/yasodak/exp/{shot}/{key}.json" + os.makedirs(f"/home/yasodak/exp/{shot}", exist_ok=True) + if os.path.exists(cp): + with open(cp) as f: + d = json.load(f) + return np.array(d['t']), np.array(d['y']) + cmd = (f"python3 -c \"" + f"import MDSplus as mds; c=mds.Connection('atlas.gat.com');" + f"c.openTree('d3d',{shot});" + f"s=c.get('ptdata2(\\\"ECE{ch:02d}\\\",{shot})');" + f"t=c.get('dim_of(ptdata2(\\\"ECE{ch:02d}\\\",{shot}))');" + f"import json; print(json.dumps({{'t':list(map(float,t)),'y':list(map(float,s))}}))\"") + r = subprocess.run(cmd, shell=True, capture_output=True, text=True) + if r.returncode != 0 or not r.stdout.strip(): + return None, None + d = json.loads(r.stdout.strip()) + with open(cp, 'w') as f: + json.dump(d, f) + return np.array(d['t']), np.array(d['y']) + +def fetch_n1rms(shot): + cp = f"/home/yasodak/exp/{shot}/mds_mhd_MHD__TOP_MIRNOV_N1RMS5.json" + with open(cp) as f: + d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: + t = t * 1e3 + return t, y + +# ── Figure 1: full flat-top overview for each shot ──────────────────────────── +for shot in SHOTS: + print(f"\nShot {shot}") + ev = LONG_EVENT[shot] + + ece_data = {} + for ch in CHANNELS: + t, y = fetch_ece(shot, ch) + if t is not None: + ece_data[ch] = (t, y) + print(f" ECE{ch:02d}: {len(t)} pts") + + t_n1, n1 = fetch_n1rms(shot) + + fig, axes = plt.subplots(len(CHANNELS)+1, 1, figsize=(14, 2*(len(CHANNELS)+1)), + sharex=True, gridspec_kw={'hspace': 0.05}) + fig.suptitle(f'Shot {shot}: ECE channels (core→edge) vs N1RMS', fontsize=11) + + # N1RMS top + ax = axes[0] + mk = (t_n1 >= 2000) & (t_n1 <= 5500) + ax.plot(t_n1[mk], n1[mk], 'k', lw=0.7) + ax.axhline(12, color='red', lw=0.7, ls='--') + ax.axvspan(*ev, color='tab:red', alpha=0.1) + ax.set_ylabel('N1RMS\n(G)', fontsize=7) + + cmap = plt.cm.plasma + for i, (ch, lbl) in enumerate(zip(CHANNELS[::-1], LABELS[::-1])): # inner→outer + ax = axes[i+1] + if ch in ece_data: + t, y = ece_data[ch] + mk = (t >= 2000) & (t <= 5500) + # downsample to 1ms for overview + step = max(1, int(5 / np.median(np.diff(t[mk])))) + ax.plot(t[mk][::step], y[mk][::step], lw=0.5, + color=cmap(i / len(CHANNELS))) + ax.axvspan(*ev, color='tab:red', alpha=0.1) + ax.set_ylabel(f'ECE{ch}\n{lbl}', fontsize=6) + ax.tick_params(labelsize=6) + + axes[-1].set_xlabel('Time (ms)', fontsize=8) + axes[-1].set_xlim(2000, 5500) + + out = f'figures/ece_overview_{shot}.png' + fig.savefig(out, dpi=120, bbox_inches='tight') + print(f" Saved {out}") + plt.close(fig) + + # ── Figure 2: zoom windows ──────────────────────────────────────────────── + for zi, (tw0, tw1, tt0, tt1) in enumerate(ZOOMS[shot]): + fig2, axes2 = plt.subplots(len(CHANNELS)+1, 1, + figsize=(12, 2*(len(CHANNELS)+1)), + sharex=True, gridspec_kw={'hspace': 0.04}) + fig2.suptitle(f'Shot {shot}: ECE zoom {tw0}–{tw1} ms ' + f'(mode onset {ev[0]} ms, yellow = tight window)', fontsize=10) + + ax = axes2[0] + mk = (t_n1 >= tw0) & (t_n1 <= tw1) + ax.plot(t_n1[mk], n1[mk], 'k', lw=0.8) + ax.axhline(12, color='red', lw=0.8, ls='--') + ax.axvspan(*ev, color='tab:red', alpha=0.15) + ax.axvspan(tt0, tt1, color='yellow', alpha=0.25) + ax.set_ylabel('N1RMS\n(G)', fontsize=7) + + for i, (ch, lbl) in enumerate(zip(CHANNELS[::-1], LABELS[::-1])): + ax2 = axes2[i+1] + if ch in ece_data: + t_e, y_e = ece_data[ch] + mk = (t_e >= tw0) & (t_e <= tw1) + ax2.plot(t_e[mk], y_e[mk], lw=0.6, + color=cmap(i / len(CHANNELS))) + mk2 = (t_e >= tt0) & (t_e <= tt1) + ax2.plot(t_e[mk2], y_e[mk2], lw=1.5, + color=cmap(i / len(CHANNELS))) + ax2.axvspan(*ev, color='tab:red', alpha=0.12) + ax2.axvspan(tt0, tt1, color='yellow', alpha=0.2) + ax2.set_ylabel(f'ECE{ch}\n{lbl}', fontsize=6) + ax2.tick_params(labelsize=6) + + axes2[-1].set_xlabel('Time (ms)', fontsize=8) + axes2[-1].set_xlim(tw0, tw1) + + out2 = f'figures/ece_zoom_{shot}_w{zi}.png' + fig2.savefig(out2, dpi=130, bbox_inches='tight') + print(f" Saved {out2}") + plt.close(fig2) + +print("\nDone.") diff --git a/src/tokeye/modespec/classic/generate_modes.py b/src/tokeye/modespec/classic/generate_modes.py new file mode 100644 index 0000000..b286dc0 --- /dev/null +++ b/src/tokeye/modespec/classic/generate_modes.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python +""" +generate_modes.py — turn a config into a list of detected MHD modes. + +Reads a YAML config that lists shots + analysis parameters, runs the modespec +pipeline on each shot, detects coherent mode "events", and writes one CSV per +shot to `/_modes.csv`. Optionally saves a four-panel +spectrogram figure per shot via plot_modespec(). + +Usage:: + + pixi run modes modes.yaml + # or + python generate_modes.py modes.yaml + +A "mode event" is a contiguous stretch of time where a single toroidal mode +number n is the dominant, statistically-significant mode. Significance per +(time, frequency) bin is: n_dominant == n AND coherence >= threshold +AND mode_amp[n] >= amp_min_G. The threshold defaults to the 95% coherence +confidence level (result['c95']) computed by mode_spectrogram(). +""" + +import sys +import csv +import argparse +from pathlib import Path + +import yaml +import numpy as np + +# Headless backend BEFORE modespec imports pyplot (modespec.py imports it at module load). +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from .modespec import fetch_mirnov, mode_spectrogram, plot_modespec + + +# ── Config handling ────────────────────────────────────────────────────────────── + +# Per-shot parameters and their defaults. A shot entry in the config may override +# any of these; anything omitted falls back to the config's `defaults` block, then +# to these built-in values. +PARAM_DEFAULTS = { + "array": "toroidal", # 'toroidal' (14-probe) or 'poloidal' (31-probe) + "integrated": False, # False = B-dot (200 kHz), True = integrated B (~20 kHz) + "t_min_ms": None, # optional signal time crop + "t_max_ms": None, + "dt_window_ms": 4.0, # FFT window length + "overlap_frac": 0.75, # window overlap fraction + "f_min_khz": 5.0, # analysis band + "f_max_khz": 50.0, + "f_smooth_khz": 1.0, # frequency smoothing bandwidth + "n_range": [1, 5], # [n_min, n_max] mode numbers to test + # detection + "coherence_min": None, # None -> use result['c95'] + "amp_min_G": 0.5, # mode_amp floor for a bin to count (units: G if + # integrated=True, else G/s — raise it for B-dot) + "merge_gap_ms": 2.0, # bridge sub-threshold gaps shorter than this + "min_duration_ms": 5.0, # discard events shorter than this + "make_figure": True, # save _modespec.png +} + +CSV_COLUMNS = [ + "array", "mode_label", "mode_number", + "t_start_ms", "t_end_ms", "duration_ms", + "peak_freq_khz", "peak_amp_G", "mean_coherence", + "f_min_khz", "f_max_khz", "coherence_thresh", +] + + +def load_config(path): + """Load YAML config -> (global_cfg, list_of_resolved_shot_cfgs).""" + with Path(path).open() as fh: + cfg = yaml.safe_load(fh) or {} + + defaults = {**PARAM_DEFAULTS, **(cfg.get("defaults") or {})} + shots = cfg.get("shots") or [] + if not shots: + raise SystemExit(f"No 'shots' listed in {path}") + + resolved = [] + for entry in shots: + if isinstance(entry, int): # allow a bare shot number + entry = {"shot": entry} + if "shot" not in entry: + raise SystemExit(f"Shot entry missing 'shot': {entry}") + merged = {**defaults, **entry} + resolved.append(merged) + + global_cfg = { + "output_dir": cfg.get("output_dir", "mode_analysis"), + "atlas": cfg.get("atlas", "atlas.gat.com"), + } + return global_cfg, resolved + + +# ── Detection ──────────────────────────────────────────────────────────────────── + +def _contiguous_runs(mask): + """Yield (start, end) inclusive index pairs for each run of True in a 1-D bool array.""" + idx = np.flatnonzero(mask) + if idx.size == 0: + return + breaks = np.flatnonzero(np.diff(idx) > 1) + starts = np.r_[idx[0], idx[breaks + 1]] + ends = np.r_[idx[breaks], idx[-1]] + for s, e in zip(starts, ends): + yield int(s), int(e) + + +def _fill_gaps(mask, max_gap): + """Bridge runs of False shorter than `max_gap` windows (morphological closing).""" + if max_gap < 1: + return mask + out = mask.copy() + inside = np.flatnonzero(mask) + if inside.size == 0: + return out + for s, e in _contiguous_runs(~mask): + if s > inside[0] and e < inside[-1] and (e - s + 1) <= max_gap: + out[s:e + 1] = True # gap is enclosed by True on both sides + return out + + +def detect_modes(result, cfg): + """Return a list of mode-event dicts from a mode_spectrogram() result.""" + t = result["t_win_ms"] + f = result["freq_khz"] + nd = result["n_dominant"] + coh = result["coherence"] + mode_amp = result["mode_amp"] + n_lo, n_hi = result["n_range"] + + thresh = cfg["coherence_min"] + if thresh is None: + thresh = result["c95"] + amp_min = cfg["amp_min_G"] + min_dur = cfg["min_duration_ms"] + + step_ms = float(np.mean(np.diff(t))) if t.size > 1 else 1.0 + max_gap = int(round(cfg["merge_gap_ms"] / step_ms)) + + events = [] + for n in range(int(n_lo), int(n_hi) + 1): + if n == 0: + continue # n=0 is axisymmetric, not a rotating mode of interest + amp_n = mode_amp[n] + sig = (nd == n) & (coh >= thresh) & (amp_n >= amp_min) # (n_win, n_freq) + present = _fill_gaps(sig.any(axis=1), max_gap) # (n_win,) + for i0, i1 in _contiguous_runs(present): + t0, t1 = float(t[i0]), float(t[i1]) + if (t1 - t0) < min_dur: + continue + # Peak (amplitude) bin among significant bins inside the event window. + sub = np.where(sig[i0:i1 + 1], amp_n[i0:i1 + 1], -np.inf) + iw, jf = np.unravel_index(np.argmax(sub), sub.shape) + events.append({ + "mode_number": n, + "t_start_ms": round(t0, 3), + "t_end_ms": round(t1, 3), + "duration_ms": round(t1 - t0, 3), + "peak_freq_khz": round(float(f[jf]), 4), + "peak_amp_G": round(float(amp_n[i0 + iw, jf]), 5), + "mean_coherence": round(float(coh[i0:i1 + 1][sig[i0:i1 + 1]].mean()), 4), + "coherence_thresh": round(float(thresh), 4), + }) + return events + + +# ── Per-shot driver ────────────────────────────────────────────────────────────── + +def process_shot(cfg, atlas, output_dir): + """Run the pipeline for one resolved shot config; return list of CSV rows.""" + shot = cfg["shot"] + array = cfg["array"] + mode_label = "n" if array == "toroidal" else "m" + print(f"[{shot}] fetching {array} array ...") + + signals, t_ms, angles, names = fetch_mirnov( + shot, array=array, integrated=cfg["integrated"], atlas=atlas, + t_min_ms=cfg["t_min_ms"], t_max_ms=cfg["t_max_ms"], + ) + + result = mode_spectrogram( + signals, t_ms, angles, + dt_window_ms=cfg["dt_window_ms"], overlap_frac=cfg["overlap_frac"], + f_min_khz=cfg["f_min_khz"], f_max_khz=cfg["f_max_khz"], + f_smooth_khz=cfg["f_smooth_khz"], n_range=tuple(cfg["n_range"]), + ) + + events = detect_modes(result, cfg) + print(f"[{shot}] detected {len(events)} mode event(s) " + f"(coh>={result['c95']:.3f}, {cfg['f_min_khz']}-{cfg['f_max_khz']} kHz)") + + if cfg["make_figure"]: + try: + fig = plot_modespec(result, shot=shot, mode_label=mode_label) + fig_path = output_dir / f"{shot}_modespec.png" + fig.savefig(fig_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"[{shot}] figure -> {fig_path}") + except Exception as exc: # never lose CSV rows over a plot + print(f"[{shot}] WARNING: figure failed: {exc}", file=sys.stderr) + + rows = [] + for ev in events: + rows.append({ + "array": array, + "mode_label": mode_label, + "f_min_khz": cfg["f_min_khz"], + "f_max_khz": cfg["f_max_khz"], + **ev, + }) + + csv_path = output_dir / f"{shot}_modes.csv" + with csv_path.open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=CSV_COLUMNS) + writer.writeheader() + writer.writerows(rows) + print(f"[{shot}] wrote {len(rows)} mode(s) -> {csv_path}") + return len(rows) + + +def run_config(config_path): + """Process every shot in a YAML config. Returns the number of failed shots.""" + global_cfg, shot_cfgs = load_config(config_path) + output_dir = Path(global_cfg["output_dir"]) + atlas = global_cfg["atlas"] + output_dir.mkdir(parents=True, exist_ok=True) + + total = ok = 0 + for cfg in shot_cfgs: + try: + total += process_shot(cfg, atlas, output_dir) # writes _modes.csv + ok += 1 + except Exception as exc: # one bad shot shouldn't kill the run + print(f"[{cfg['shot']}] ERROR: {exc}", file=sys.stderr) + + print(f"\nWrote {total} mode(s) across {ok}/{len(shot_cfgs)} shot(s) " + f"-> {output_dir}/_modes.csv") + return len(shot_cfgs) - ok + + +def main(): + ap = argparse.ArgumentParser(description="Generate a list of MHD modes from a config.") + ap.add_argument("config", nargs="?", default="modes.yaml", + help="YAML config (default: modes.yaml)") + args = ap.parse_args() + return run_config(args.config) + + +if __name__ == "__main__": + main() diff --git a/src/tokeye/modespec/classic/modes.yaml b/src/tokeye/modespec/classic/modes.yaml new file mode 100644 index 0000000..50315fb --- /dev/null +++ b/src/tokeye/modespec/classic/modes.yaml @@ -0,0 +1,59 @@ +# modes.yaml — declare which shots to analyze and how. +# +# `generate_modes.py` reads this file, runs the modespec pipeline on each shot, +# and writes detected mode events to /modes.csv (+ a figure per shot). +# +# pixi run modes modes.yaml +# +# Each shot inherits everything in `defaults`; list a key under a shot to override +# it. To analyze a shot in two frequency bands, list it twice with different +# f_min_khz/f_max_khz. + +output_dir: mode_analysis # outputs land here (already gitignored) +atlas: atlas.gat.com # MDSplus server + +defaults: + array: toroidal # toroidal (14-probe n-analysis) | poloidal (31-probe m-analysis) + integrated: false # false = B-dot signals (200 kHz) | true = integrated B (~20 kHz) + t_min_ms: null # optional crop of the signal time range (null = full shot) + t_max_ms: null + + # spectrogram + dt_window_ms: 4.0 # FFT window length + overlap_frac: 0.75 # window overlap fraction (0-1) + f_min_khz: 5.0 # analysis band, low edge + f_max_khz: 50.0 # analysis band, high edge + f_smooth_khz: 1.0 # frequency smoothing bandwidth (0 disables) + n_range: [1, 5] # [n_min, n_max] toroidal mode numbers to test + + # detection (what counts as a "mode") + coherence_min: null # null -> use the 95% confidence level (result['c95']) + amp_min_G: 0.5 # min mode amplitude floor; units are G if integrated=true, + # else G/s (B-dot) where values are large -> effectively + # coherence-only. Raise it to thin out B-dot detections. + merge_gap_ms: 2.0 # bridge sub-threshold gaps shorter than this (de-fragment) + min_duration_ms: 5.0 # discard events shorter than this + make_figure: true # also save _modespec.png + +shots: + - shot: 170008 + - shot: 170660 + - shot: 170670 + - shot: 170672 + - shot: 170677 + - shot: 170679 + - shot: 170796 + - shot: 175987 + - shot: 176054 + - shot: 178631 + - shot: 180634 + - shot: 184847 + - shot: 184859 + - shot: 184902 + - shot: 184964 + - shot: 185781 + - shot: 190904 + - shot: 193273 + - shot: 193277 + - shot: 193280 + - shot: 193281 diff --git a/src/tokeye/modespec/classic/modespec.py b/src/tokeye/modespec/classic/modespec.py new file mode 100644 index 0000000..7c39d45 --- /dev/null +++ b/src/tokeye/modespec/classic/modespec.py @@ -0,0 +1,1724 @@ +""" +modespec.py — Python implementation of DIII-D Mirnov mode analysis. + +Replicates the core algorithm from /fusion/usc/src/idl/modespec_auto/modespec.pro. +Fetches toroidal Mirnov array (Bp_probes_R0) or poloidal array (Bp_probes_322) +from MDSplus and computes: + + - Power spectrogram (FFT in sliding windows) + - Toroidal mode number spectrogram (matched-filter fit) + - RMS amplitude per mode number vs time + - Single time-slice coherence-weighted phase fit + - Multi-harmonic least-squares fit (IDL slice_fit equivalent) + +Probe geometry from: + /fusion/projects/diagnostics/magnetics/data/coords/all_mag + +Signal naming convention: + B-dot signals (200 kHz): MPI66M020D, MPI66M067D, ... (integrated=False) + Integrated B signals (~20 kHz): MPI66M020, MPI66M067, ... (integrated=True) + +Usage (notebook):: + + from modespec import fetch_mirnov, mode_spectrogram, plot_modespec + + signals, t_ms, phi_tor, names = fetch_mirnov(shot, integrated=True) + result = mode_spectrogram(signals, t_ms, phi_tor, + dt_window_ms=10.0, f_smooth_khz=1.0, + f_min_khz=1, f_max_khz=8, n_range=(-5, 5)) + fig = plot_modespec(result, shot, n1rms_dict=n1rms[shot]) +""" + +import numpy as np +import matplotlib.pyplot as plt + +# numpy 2.x removed np.trapz (renamed to np.trapezoid); restore the old name so +# this module works under both numpy 1.x (cluster) and 2.x (pixi/conda-forge). +if not hasattr(np, "trapz"): + np.trapz = np.trapezoid + +# ── Probe geometry ───────────────────────────────────────────────────────────── +# Bp_probes_R0: midplane toroidal array, 14 probes +# (name, phi_deg) from all_mag. Re-centred: probes > 315° get phi - 360°. +# B-dot names end in D; integrated names are the same without the trailing D. +TOR_PROBES = [ + ('MPI66M020D', 19.5), + ('MPI66M067D', 67.5), + ('MPI66M097D', 97.4), + ('MPI66M127D', 127.9), + ('MPI66M132D', 132.5), + ('MPI66M137D', 137.4), + ('MPI66M157D', 157.6), + ('MPI66M200D', 199.7), + ('MPI66M247D', 246.4), + ('MPI66M277D', 277.5), + ('MPI66M307D', 307.0), + ('MPI66M312D', 312.4), + ('MPI66M322D', 317.4), # actual location per all_mag, not 322° + ('MPI66M340D', 339.7), +] +TOR_PROBE_NAMES = [p[0] for p in TOR_PROBES] +TOR_PHI_RAW = np.array([p[1] for p in TOR_PROBES]) +TOR_PHI_DEG = np.where(TOR_PHI_RAW > 315, TOR_PHI_RAW - 360, TOR_PHI_RAW) +_sort_idx = np.argsort(TOR_PHI_DEG) +TOR_PROBE_NAMES = [TOR_PROBE_NAMES[i] for i in _sort_idx] +TOR_PHI_DEG = TOR_PHI_DEG[_sort_idx] + +# Bp_probes_322: poloidal array at phi ≈ 322°, 31 probes +# (name, R[m], Z[m]); theta = atan2(Z, R - R0) +POL_PROBES_RAW = [ + ('MPI11M322D', 0.973, -0.002), + ('MPI1A322D', 0.974, 0.182), + ('MPI2A322D', 0.974, 0.512), + ('MPI3A322D', 0.975, 0.850), + ('MPI4A322D', 0.972, 1.161), + ('MPI5A322D', 1.051, 1.330), + ('MPI8A322D', 1.219, 1.406), + ('MPI89A322D', 1.402, 1.407), + ('MPI9A322D', 1.584, 1.408), + ('MPI79FA322D',1.783, 1.323), + ('MPI79NA322D',1.924, 1.206), + ('MPI7FA322D', 2.067, 1.090), + ('MPI7NA322D', 2.219, 0.870), + ('MPI67A322D', 2.270, 0.746), + ('MPI6FA322D', 2.319, 0.623), + ('MPI6NA322D', 2.416, 0.249), + ('MPI66M322D', 2.418, -0.001), + ('MPI1B322D', 0.974, -0.187), + ('MPI2B322D', 0.975, -0.512), + ('MPI3B322D', 0.974, -0.854), + ('MPI4B322D', 0.972, -1.159), + ('MPI5B322D', 1.048, -1.330), + ('MPI8B322D', 1.254, -1.405), + ('MPI89B322D', 1.477, -1.406), + ('MPI9B322D', 1.699, -1.406), + ('MPI79B322D', 1.894, -1.333), + ('MPI7FB322D', 2.085, -1.102), + ('MPI7NB322D', 2.212, -0.873), + ('MPI67B322D', 2.263, -0.749), + ('MPI6FB322D', 2.315, -0.624), + ('MPI6NB322D', 2.416, -0.244), +] +POL_R0 = 1.69 +POL_PROBE_NAMES = [p[0] for p in POL_PROBES_RAW] +_R = np.array([p[1] for p in POL_PROBES_RAW]) +_Z = np.array([p[2] for p in POL_PROBES_RAW]) +POL_THETA_DEG = np.degrees(np.arctan2(_Z, _R - POL_R0)) +POL_THETA_DEG = np.where(POL_THETA_DEG < -5, POL_THETA_DEG + 360, POL_THETA_DEG) +_psort = np.argsort(POL_THETA_DEG) +POL_PROBE_NAMES = [POL_PROBE_NAMES[i] for i in _psort] +POL_THETA_DEG = POL_THETA_DEG[_psort] + + +# ── ECE radial positions ─────────────────────────────────────────────────────── + +def _ece_freq_ghz(shot, n_ch=40): + """ECE channel frequencies [GHz] for DIII-D (from getecefreq.py formula).""" + freqs = np.zeros(n_ch) + if shot > 178000: + freqs[:16] = 82.5 + np.arange(16) # 82.5–97.5 GHz + freqs[16:32] = 98.5 + np.arange(16) # 98.5–113.5 GHz + freqs[32:40] = 115.5 + 2 * np.arange(8) # 115.5–129.5 GHz (step 2) + elif shot > 100600: + freqs[:16] = 83.5 + np.arange(16) + freqs[16:32] = 98.5 + np.arange(16) + freqs[32:40] = 115.5 + 2 * np.arange(8) + else: + freqs[:16] = 83.5 + np.arange(16) + freqs[16:32] = 98.5 + np.arange(16) + return freqs + + +def ece_channel_radius(shot, bt0_T=None, r_axis_m=None): + """ + Compute major radius R [m] for each fixed ECE channel. + + Uses second-harmonic emission: f_ECE = 2 × f_ce = 56.0 × B_T [GHz/T]. + With B_T(R) = B_T0 × R_axis / R: R = 56.0 × B_T0 × R_axis / f_ECE. + + Parameters + ---------- + shot : DIII-D shot number (for frequency table lookup) + bt0_T : toroidal field at magnetic axis [T]. Default: 1.76 T (typical H-mode) + r_axis_m : major radius of magnetic axis [m]. Default: 1.69 m + + Returns + ------- + freq_ghz : (n_ch,) ECE frequency per channel [GHz] + R_ece : (n_ch,) corresponding major radius [m] + """ + if bt0_T is None: + bt0_T = 1.76 + if r_axis_m is None: + r_axis_m = 1.69 + freq_ghz = _ece_freq_ghz(shot) + R_ece = 56.0 * bt0_T * r_axis_m / freq_ghz + return freq_ghz, R_ece + + +# ── MDSplus fetch ────────────────────────────────────────────────────────────── + +def fetch_mirnov(shot, array='toroidal', integrated=False, + atlas='atlas.gat.com', t_min_ms=None, t_max_ms=None): + """ + Fetch Mirnov array signals from MDSplus. + + Parameters + ---------- + shot : int + array : 'toroidal' → 14-probe midplane toroidal array + 'poloidal' → 31-probe 322° poloidal array + integrated : bool + False (default) → B-dot signals (names end in D), 200 kHz + True → integrated B signals (no trailing D), ~20 kHz + atlas : MDSplus server + t_min_ms, t_max_ms : optional time crop in ms + + Returns + ------- + signals : ndarray (n_probes, n_t) [G/s for B-dot; G for integrated] + t_ms : ndarray (n_t,) time axis [ms] + angles : ndarray (n_probes,) phi_tor or theta_pol [deg] + names : list of str PTDATA signal names used + """ + try: + import MDSplus as mds + except ImportError: + raise RuntimeError( + 'MDSplus is not installed. It ships on the GA cluster and on ' + 'conda-forge (conda install -c conda-forge mdsplus); fetching ' + 'DIII-D data also requires network access to atlas.gat.com.' + ) + + if array == 'toroidal': + probe_names = list(TOR_PROBE_NAMES) + angles = TOR_PHI_DEG.copy() + else: + probe_names = list(POL_PROBE_NAMES) + angles = POL_THETA_DEG.copy() + + if integrated: + # Strip trailing 'D' for integrated B signals + probe_names = [n[:-1] if n.endswith('D') else n for n in probe_names] + + conn = mds.Connection(atlas) + conn.openTree('D3D', shot) + + signals = [] + good_names = [] + good_angles= [] + t_ref = None # full time axis from first successful probe + + for name, phi in zip(probe_names, angles): + bdot_fallback = False + try: + d = np.array(conn.get(f'PTDATA("{name}",{shot})').data(), dtype=float) + t = np.array(conn.get(f'DIM_OF(PTDATA("{name}",{shot}))').data(), dtype=float) + except Exception: + if integrated and not name.endswith('D'): + # Integrated version absent — fall back to B-dot and integrate + try: + bdot_name = name + 'D' + d = np.array(conn.get(f'PTDATA("{bdot_name}",{shot})').data(), dtype=float) + t = np.array(conn.get(f'DIM_OF(PTDATA("{bdot_name}",{shot}))').data(), dtype=float) + bdot_fallback = True + except Exception as e2: + print(f' Warning: {name} (and {name}D) failed: {e2}') + continue + else: + import sys + print(f' Warning: {name} failed', file=sys.stderr) + continue + + if t.size < 2: + continue + + if bdot_fallback: + # Numerical integration of B-dot → B (cumulative trapezoid) + dt_s = np.mean(np.diff(t)) * 1e-3 # ms → s + d = np.cumsum(d) * dt_s # ∫ B-dot dt [G] + # High-pass: subtract running mean to remove integration drift + from scipy.ndimage import uniform_filter1d + hp_pts = max(1, int(round(200.0 / (dt_s * 1e3)))) # 200 ms window + d = d - uniform_filter1d(d, size=hp_pts, mode='nearest') + name = name + 'D→∫' + + if t_ref is None: + t_ref = t.copy() + if t_min_ms is not None or t_max_ms is not None: + lo = t_min_ms if t_min_ms is not None else t[0] + hi = t_max_ms if t_max_ms is not None else t[-1] + mask = (t >= lo) & (t <= hi) + t = t[mask] + d = d[mask] + signals.append(d) + good_names.append(name) + good_angles.append(phi) + + conn.closeAllTrees() + + if not signals: + raise RuntimeError(f'No Mirnov signals fetched for shot {shot}') + + # Build common time axis from first probe's reference + if t_ref is None: + t_ref = np.linspace(0, 1, len(signals[0])) + if t_min_ms is not None or t_max_ms is not None: + lo = t_min_ms if t_min_ms is not None else t_ref[0] + hi = t_max_ms if t_max_ms is not None else t_ref[-1] + t_ref = t_ref[(t_ref >= lo) & (t_ref <= hi)] + + n_t_min = min(len(s) for s in signals) + data = np.array([s[:n_t_min] for s in signals]) # (n_probes, n_t) + t_ms = t_ref[:n_t_min] + + fs_khz = 1.0 / (float(np.mean(np.diff(t_ms))) * 1e-3) / 1e3 + sig_type = 'integrated B' if integrated else 'B-dot' + print(f' Fetched {len(good_names)}/{len(probe_names)} {array} probes ' + f'({sig_type}), t=[{t_ms[0]:.0f},{t_ms[-1]:.0f}] ms, fs={fs_khz:.0f} kHz') + + return data, t_ms, np.array(good_angles), good_names + + +def fetch_surfmn(shot, atlas='atlas.gat.com', t_min_ms=None, t_max_ms=None, + modes=None): + """ + Fetch pre-computed SURFMN (m,n) mode amplitudes and radial locations + from \\MHD::TOP.SURFMN.OUTPUT.ISLTABLE. + + B_M_N[n_idx, m_idx, t] where n_idx = n-1 (0-based) and m_idx = m (direct). + ~20 ms cadence, full-shot coverage. + + Parameters + ---------- + shot : int + atlas : MDSplus server + t_min_ms, t_max_ms : optional time crop [ms] + modes : list of (m, n) tuples to return; default [(2,1),(3,1),(4,1),(3,2)] + + Returns + ------- + dict with keys: + 't_ms' : (n_t,) time axis [ms] + 'amp' : {(m,n): (n_t,) amplitude array} + 'rho' : {(m,n): (n_t,) radial location (rho_N)} + """ + if modes is None: + modes = [(2, 1), (3, 1), (4, 1), (3, 2)] + + try: + import MDSplus as mds + except ImportError: + raise RuntimeError( + 'MDSplus is not installed. It ships on the GA cluster and on ' + 'conda-forge (conda install -c conda-forge mdsplus); fetching ' + 'DIII-D data also requires network access to atlas.gat.com.' + ) + + conn = mds.Connection(atlas) + conn.openTree('MHD', shot) + base = '\\MHD::TOP.SURFMN.OUTPUT.ISLTABLE' + + bmn = np.array(conn.get(f'{base}:B_M_N').data()) # (n_max, m_max, n_t) + loc = np.array(conn.get(f'{base}:LOC_M_N').data()) + t = np.array(conn.get(f'DIM_OF({base}:B_M_N)').data()) # ms + conn.closeAllTrees() + + mask = np.ones(len(t), dtype=bool) + if t_min_ms is not None: + mask &= t >= t_min_ms + if t_max_ms is not None: + mask &= t <= t_max_ms + t = t[mask] + + amp_out = {} + rho_out = {} + for (m, n) in modes: + n_idx = n - 1 + m_idx = m + if n_idx < bmn.shape[0] and m_idx < bmn.shape[1]: + amp_out[(m, n)] = bmn[n_idx, m_idx, mask] + rho_out[(m, n)] = loc[n_idx, m_idx, mask] + else: + amp_out[(m, n)] = np.zeros(mask.sum()) + rho_out[(m, n)] = np.zeros(mask.sum()) + + print(f' SURFMN shot {shot}: {len(t)} time pts ' + f't=[{t[0]:.0f},{t[-1]:.0f}] ms, modes={modes}') + return {'t_ms': t, 'amp': amp_out, 'rho': rho_out} + + +def plot_surfmn(results, shot_labels=None, onset_ms=None, n1rms_dict=None, + modes=None, figsize=(11, 6)): + """ + Plot SURFMN mode amplitudes and radial locations for one or more shots. + + Parameters + ---------- + results : dict {shot: fetch_surfmn output} or single fetch_surfmn output + shot_labels: dict {shot: label string} + onset_ms : dict {shot: onset time [ms]} or scalar + n1rms_dict : dict {shot: n1rms fetch_overview dict} — adds N1RMS panel if provided + modes : list of (m,n) to plot; default all modes in first result + """ + if not isinstance(results, dict) or 't_ms' in results: + results = {0: results} + if shot_labels is None: + shot_labels = {0: ''} + + shots = list(results.keys()) + if modes is None: + modes = list(next(iter(results.values()))['amp'].keys()) + if shot_labels is None: + shot_labels = {s: str(s) for s in shots} + + colors_mode = {(2,1): 'C3', (3,1): 'C1', (4,1): 'C4', + (3,2): 'C2', (1,1): 'C5'} + ls_shot = ['-', '--', ':'] + + n_panels = 3 if n1rms_dict else 2 + fig, axes = plt.subplots(n_panels, 1, figsize=figsize, sharex=True) + ax_amp, ax_rho = axes[0], axes[1] + ax_rms = axes[2] if n_panels == 3 else None + + for i_s, shot in enumerate(shots): + res = results[shot] + t = res['t_ms'] + ls = ls_shot[i_s % len(ls_shot)] + lbl_shot = shot_labels.get(shot, str(shot)) + + for (m, n) in modes: + amp = res['amp'].get((m, n), np.zeros_like(t)) + rho = res['rho'].get((m, n), np.zeros_like(t)) + col = colors_mode.get((m, n), f'C{m}') + label = f'{lbl_shot} {m}/{n}' if len(shots) > 1 else f'm/n={m}/{n}' + ax_amp.plot(t, amp, color=col, ls=ls, lw=1.4, label=label) + mask = rho > 0 + if mask.any(): + ax_rho.plot(t[mask], rho[mask], color=col, ls=ls, lw=1.4) + + ax_amp.set_ylabel('SURFMN amplitude [a.u.]') + ax_amp.legend(fontsize=7, ncol=2) + ax_amp.grid(True, alpha=0.3) + + ax_rho.set_ylabel(r'$\rho_N$ location') + ax_rho.set_ylim(0, 1) + ax_rho.axhline(0.59, color='gray', ls=':', lw=0.8, alpha=0.5) + ax_rho.grid(True, alpha=0.3) + + if ax_rms is not None: + for i_s, shot in enumerate(shots): + nd = n1rms_dict.get(shot, {}) + t_r = np.array(nd.get('time', [])) + v_r = np.array(nd.get('n1rms5', nd.get('n1rms', []))) + if t_r.size and v_r.size: + ax_rms.plot(t_r, v_r, lw=1.2, ls=ls_shot[i_s % 3], + label=shot_labels.get(shot, str(shot))) + ax_rms.set_ylabel('N1RMS [G]') + if ax_rms.get_legend_handles_labels()[0]: + ax_rms.legend(fontsize=7) + ax_rms.grid(True, alpha=0.3) + ax_rms.set_xlabel('Time [ms]') + else: + ax_rho.set_xlabel('Time [ms]') + + # Onset markers + if onset_ms is not None: + ons = onset_ms if isinstance(onset_ms, dict) else {shots[0]: onset_ms} + for i_s, shot in enumerate(shots): + t_on = ons.get(shot) + if t_on is not None: + for ax in axes: + ax.axvline(t_on, color='k', ls='--', lw=0.9, alpha=0.6) + + shot_str = ' / '.join(str(s) for s in shots) + ax_amp.set_title(f'Shot {shot_str} — SURFMN mode amplitudes', fontsize=10) + fig.tight_layout() + return fig + + +def fetch_ece(shot, atlas='atlas.gat.com', n_channels=40, + t_min_ms=None, t_max_ms=None, bt0_T=None, r_axis_m=None, + fast=True): + """ + Fetch ECE Te channels from the 'ece' MDSplus tree. + + fast=True (default): \\TECEF01…\\TECEF{n} 500 kHz — for fluctuation/burst analysis + fast=False : \\TECE01…\\TECE{n} ~5 kHz — for time-averaged Te profiles + + Radial positions computed from B_T(R) = B_T0 × R_axis / R at second harmonic. + + Parameters + ---------- + shot : int + atlas : MDSplus server + n_channels : number of fixed-frequency channels to fetch (default 40) + t_min_ms, t_max_ms : optional time crop [ms] + bt0_T : B_T at magnetic axis [T] for R_ece calculation + r_axis_m : R of magnetic axis [m] + fast : bool — True → TECEF (500 kHz); False → TECE (~5 kHz) + + Returns + ------- + signals : ndarray (n_ch, n_t) Te [keV] + t_ms : ndarray (n_t,) time axis [ms] + R_ece : ndarray (n_ch,) major radius per channel [m] + freq_ghz: ndarray (n_ch,) ECE frequency per channel [GHz] + ch_idx : list of int channel indices that were successfully fetched + """ + try: + import MDSplus as mds + except ImportError: + raise RuntimeError( + 'MDSplus is not installed. It ships on the GA cluster and on ' + 'conda-forge (conda install -c conda-forge mdsplus); fetching ' + 'DIII-D data also requires network access to atlas.gat.com.' + ) + + freq_ghz, R_ece = ece_channel_radius(shot, bt0_T=bt0_T, r_axis_m=r_axis_m) + + conn = mds.Connection(atlas) + conn.openTree('ece', shot) + + signals = [] + good_ch = [] + good_R = [] + good_f = [] + t_ref = None + + prefix = 'TECEF' if fast else 'TECE' + for i in range(1, n_channels + 1): + name = f'\\{prefix}{i:02d}' + try: + d = np.array(conn.get(name).data(), dtype=float) + t = np.array(conn.get(f'dim_of({name})').data(), dtype=float) + if t.size < 2: + continue + if t_ref is None: + t_ref = t.copy() + if t_min_ms is not None or t_max_ms is not None: + lo = t_min_ms if t_min_ms is not None else t[0] + hi = t_max_ms if t_max_ms is not None else t[-1] + mask = (t >= lo) & (t <= hi) + t = t[mask] + d = d[mask] + signals.append(d) + good_ch.append(i) + good_R.append(R_ece[i - 1]) + good_f.append(freq_ghz[i - 1]) + except Exception as e: + print(f' Warning: ECE ch{i:02d} failed: {e}') + + conn.closeAllTrees() + + if not signals: + raise RuntimeError(f'No ECE channels fetched for shot {shot}') + + if t_ref is None: + t_ref = np.linspace(0, 1, len(signals[0])) + if t_min_ms is not None or t_max_ms is not None: + lo = t_min_ms if t_min_ms is not None else t_ref[0] + hi = t_max_ms if t_max_ms is not None else t_ref[-1] + t_ref = t_ref[(t_ref >= lo) & (t_ref <= hi)] + + n_t = min(len(s) for s in signals) + data = np.array([s[:n_t] for s in signals]) + t_ms = t_ref[:n_t] + fs_khz = 1.0 / (float(np.mean(np.diff(t_ms))) * 1e-3) / 1e3 + + print(f' Fetched {len(good_ch)}/{n_channels} ECE channels, ' + f't=[{t_ms[0]:.0f},{t_ms[-1]:.0f}] ms, fs={fs_khz:.0f} kHz, ' + f'R=[{min(good_R):.3f},{max(good_R):.3f}] m') + + return (data, t_ms, + np.array(good_R), + np.array(good_f), + good_ch) + + +def ece_mode_location(ece_signals, t_ms, R_ece, + f_mode_khz, t_start_ms, t_end_ms, + df_band_khz=3.0, overview=None, + ref_signal=None, ref_t_ms=None): + """ + Locate a rotating MHD mode radially from ECE Te oscillation amplitude. + + Two modes (controlled by ref_signal): + + ref_signal=None [default — bandpass RMS] + Bandpass each ECE channel at f_mode_khz ± df_band_khz/2, compute + δTe/Te RMS over [t_start_ms, t_end_ms]. + + ref_signal provided [coherence mode] + Compute the mean-squared coherence γ²(f_mode) between each ECE channel + and the reference (e.g. N1RMS interpolated to the ECE timebase). + Coherence is dimensionless (0–1) and insensitive to absolute Te level, + removing the optical-depth bias that makes core channels dominate in + the RMS method. ref_t_ms must also be provided (timebase of ref_signal). + + Parameters + ---------- + ece_signals : (n_ch, n_t) Te [keV] + t_ms : (n_t,) time [ms] + R_ece : (n_ch,) major radius per channel [m] + f_mode_khz : float, mode rotation frequency [kHz] + t_start_ms, t_end_ms : time window for amplitude analysis [ms] + df_band_khz : full bandpass width [kHz] + overview : dict from overview_ntm.pkl for rho estimate (optional) + ref_signal : (n_ref,) reference signal (e.g. N1RMS) — enables coherence mode + ref_t_ms : (n_ref,) timebase of ref_signal [ms] + + Returns + ------- + dict with: + R_ece : (n_ch,) major radius array [m] + rms_vs_R : (n_ch,) metric vs R (δTe/Te in RMS mode; γ² in coherence mode) + R_peak : float, R at maximum metric [m] + ch_peak : int, channel index (0-based) at peak + rho_peak : float or None, normalised rho at peak (outboard) + f_mode_khz : float + t_start_ms, t_end_ms : float + method : 'rms' or 'coherence' + """ + from scipy.signal import butter, filtfilt, coherence as sp_coherence + + dt_ms = float(np.mean(np.diff(t_ms))) + fs_hz = 1e3 / dt_ms + f_lo = max(1.0, f_mode_khz - df_band_khz / 2.0) + f_hi = f_mode_khz + df_band_khz / 2.0 + nyq = 0.5 * fs_hz + lo_n, hi_n = f_lo * 1e3 / nyq, f_hi * 1e3 / nyq + hi_n = min(hi_n, 0.999) + + b, a = butter(4, [lo_n, hi_n], btype='band') + mask = (t_ms >= t_start_ms) & (t_ms <= t_end_ms) + + use_coherence = ref_signal is not None and ref_t_ms is not None + rms_vs_R = np.zeros(len(R_ece)) + + if use_coherence: + # Interpolate reference onto ECE timebase; restrict to analysis window + ref_interp = np.interp(t_ms, ref_t_ms, ref_signal) + ref_win = ref_interp[mask] - np.mean(ref_interp[mask]) + nperseg = min(256, mask.sum() // 4) + nperseg = max(nperseg, 16) + for j in range(len(R_ece)): + sig = ece_signals[j, mask] + Te0 = float(np.median(sig)) + if Te0 < 0.05: + continue + sig_hp = sig - np.mean(sig) + try: + f_coh, coh = sp_coherence(sig_hp, ref_win, fs=fs_hz, + nperseg=nperseg) + # mean γ² over the mode frequency band + band = (f_coh >= f_lo * 1e3) & (f_coh <= f_hi * 1e3) + rms_vs_R[j] = float(np.mean(coh[band])) if band.any() else 0.0 + except Exception: + rms_vs_R[j] = 0.0 + method = 'coherence' + else: + for j in range(len(R_ece)): + sig = ece_signals[j] + Te0 = float(np.median(sig)) + sig_hp = sig - Te0 + try: + filt = filtfilt(b, a, sig_hp) + except Exception: + filt = sig_hp + abs_rms = float(np.sqrt(np.mean(filt[mask] ** 2))) + rms_vs_R[j] = abs_rms / Te0 if Te0 > 0.05 else 0.0 + method = 'rms' + + ch_peak = int(np.argmax(rms_vs_R)) + R_peak = float(R_ece[ch_peak]) + + # rho_N per channel: |R - R_axis| / a_minor (unsigned; override with EFIT in caller) + rho_peak = None + rho_vs_R = None + r_ax_ref = None + a_min_ref = None + if overview is not None: + t_mid = 0.5 * (t_start_ms + t_end_ms) + r_ax_ref = float(np.interp(t_mid, + np.array(overview.get('rmaxis', {}).get('time', [t_mid])), + np.array(overview.get('rmaxis', {}).get('data', [1.69])))) + a_min_ref = float(np.interp(t_mid, + np.array(overview.get('aminor', {}).get('time', [t_mid])), + np.array(overview.get('aminor', {}).get('data', [0.60])))) + if a_min_ref > 0: + rho_vs_R = np.abs(R_ece - r_ax_ref) / a_min_ref + rho_peak = float(rho_vs_R[ch_peak]) + + return { + 'R_ece': R_ece, + 'rms_vs_R': rms_vs_R, + 'R_peak': R_peak, + 'ch_peak': ch_peak, + 'rho_peak': rho_peak, + 'rho_vs_R': rho_vs_R, + 'r_axis_m': r_ax_ref, + 'a_minor_m': a_min_ref, + 'f_mode_khz': f_mode_khz, + 't_start_ms': t_start_ms, + 't_end_ms': t_end_ms, + 'method': method, + } + + +def plot_ece_location(loc_result, shot=None, rho_q2=None, figsize=(9, 4)): + """ + Two-panel ECE mode location plot. + + Left: RMS amplitude vs major radius R (marks peak channel) + Right: RMS amplitude vs rho_N (outboard midplane), with q=2 surface + """ + R = loc_result['R_ece'] + rms = loc_result['rms_vs_R'] + R_pk = loc_result['R_peak'] + rho = loc_result.get('rho_peak') + f0 = loc_result['f_mode_khz'] + t0 = loc_result['t_start_ms'] + t1 = loc_result['t_end_ms'] + + sort_r = np.argsort(R) + R_s = R[sort_r] + rms_s = rms[sort_r] + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize) + title = (f'Shot {shot} — ' if shot else '') + \ + f'ECE mode location f={f0:.1f} kHz t=[{t0:.0f},{t1:.0f}] ms' + fig.suptitle(title, fontsize=9) + + method = loc_result.get('method', 'rms') + y_label = 'γ² at f_mode (coherence)' if method == 'coherence' else 'δTe/Te (%)' + y_scale = 1.0 if method == 'coherence' else 1e2 + + # Left: metric vs R + ax1.plot(R_s, rms_s * y_scale, 'o-', lw=1.2, ms=4) + ax1.axvline(R_pk, color='r', ls='--', lw=1, label=f'R={R_pk:.3f} m') + ax1.set_xlabel('Major radius R (m)') + ax1.set_ylabel(y_label) + ax1.legend(fontsize=8) + ax1.grid(False) + + # Right: RMS vs rho_N (outboard midplane) + rho_arr = loc_result.get('rho_vs_R') + if rho_arr is not None: + sort_rho = np.argsort(rho_arr) + rho_s = rho_arr[sort_rho] + rms_rho_s = rms[sort_rho] + rho_pk = rho + ax2.plot(rho_s, rms_rho_s * y_scale, 'o-', lw=1.2, ms=4, color='C0') + if rho_pk is not None: + ax2.axvline(rho_pk, color='r', ls='--', lw=1, label=f'ρ_N={rho_pk:.2f}') + if rho_q2 is not None: + ax2.axvline(rho_q2, color='k', ls=':', lw=1.2, label=f'q=2 ρ={rho_q2:.2f}') + ax2.set_xlabel('ρ_N (outboard midplane)') + ax2.set_ylabel(y_label) + ax2.legend(fontsize=8) + ax2.grid(False) + else: + ax2.plot(np.arange(len(rms)), rms * y_scale, 'o-', lw=1.2, ms=4) + ax2.set_xlabel('ECE channel index') + ax2.set_ylabel('δTe/Te (%)') + ax2.grid(False) + + fig.tight_layout(rect=[0, 0, 1, 0.93]) + return fig + + +# ── EFIT q-map ──────────────────────────────────────────────────────────────── + +def fetch_efit_qmap(shot, t_ms, atlas='atlas.gat.com', tree='EFIT02ER'): + """ + Fetch 2-D q map from EFIT02ER at the time slice closest to t_ms. + + Reads PSIRZ (nt, nZ, nR), QPSI (nt, nq), SSIMAG/SSIBRY (nt,), + RMAXIS/ZMAXIS/RBBBS/ZBBBS from the EFIT02ER MDSplus tree. + + Returns + ------- + dict with keys: + R_grid : (nR,) major radius grid [m] + Z_grid : (nZ,) height grid [m] + psi_n : (nZ, nR) normalised poloidal flux at t_actual + q_2d : (nZ, nR) q value mapped onto spatial grid + q_psi : (nq,) q profile on uniform psi_N = linspace(0,1,nq) + psi_1d : (nq,) psi_N grid for q_psi + t_actual : float actual time of slice [ms] + R_axis : float R of magnetic axis [m] + Z_axis : float Z of magnetic axis [m] + R_bdy : (nb,) LCFS R [m] + Z_bdy : (nb,) LCFS Z [m] + """ + try: + import MDSplus as mds + except ImportError: + raise RuntimeError( + 'MDSplus is not installed. It ships on the GA cluster and on ' + 'conda-forge (conda install -c conda-forge mdsplus); fetching ' + 'DIII-D data also requires network access to atlas.gat.com.' + ) + from scipy.interpolate import interp1d + + base = f'\\{tree}::TOP.RESULTS.GEQDSK:' + conn = mds.Connection(atlas) + conn.openTree(tree, shot) + + psirz = np.array(conn.get(f'{base}PSIRZ').data(), dtype=float) # (nt, nZ, nR) + r_arr = np.array(conn.get(f'dim_of({base}PSIRZ, 0)').data(), dtype=float) + z_arr = np.array(conn.get(f'dim_of({base}PSIRZ, 1)').data(), dtype=float) + t_arr = np.array(conn.get(f'dim_of({base}PSIRZ, 2)').data(), dtype=float) + qpsi = np.array(conn.get(f'{base}QPSI').data(), dtype=float) # (nt, nq) + ssimag = np.array(conn.get(f'{base}SSIMAG').data(), dtype=float) # (nt,) + ssibry = np.array(conn.get(f'{base}SSIBRY').data(), dtype=float) # (nt,) + rmaxis = np.array(conn.get(f'{base}RMAXIS').data(), dtype=float) + zmaxis = np.array(conn.get(f'{base}ZMAXIS').data(), dtype=float) + rbbbs = np.array(conn.get(f'{base}RBBBS').data(), dtype=float) # (nt, nb) + zbbbs = np.array(conn.get(f'{base}ZBBBS').data(), dtype=float) + conn.closeAllTrees() + + ti = int(np.argmin(np.abs(t_arr - t_ms))) + t_actual = float(t_arr[ti]) + + psi_slice = psirz[ti] # (nZ, nR) + dpsi = float(ssibry[ti] - ssimag[ti]) + psi_n = (psi_slice - ssimag[ti]) / dpsi # (nZ, nR) 0=axis, 1=LCFS + + nq = qpsi.shape[1] + psi_1d = np.linspace(0.0, 1.0, nq) + q_row = qpsi[ti] # (nq,) from axis to boundary + q_func = interp1d(psi_1d, q_row, kind='linear', + bounds_error=False, fill_value=(q_row[0], q_row[-1])) + q_2d = q_func(np.clip(psi_n, 0.0, 1.0)) # (nZ, nR) + + print(f' EFIT {tree} t={t_actual:.0f} ms ' + f'q_axis={q_row[0]:.2f} q(psi_N=0.9)={q_func(0.9):.2f} ' + f'R_axis={rmaxis[ti]:.3f} m') + + return { + 'R_grid': r_arr, + 'Z_grid': z_arr, + 'psi_n': psi_n, + 'q_2d': q_2d, + 'q_psi': q_row, + 'psi_1d': psi_1d, + 't_actual': t_actual, + 'R_axis': float(rmaxis[ti]), + 'Z_axis': float(zmaxis[ti]), + 'R_bdy': rbbbs[ti], + 'Z_bdy': zbbbs[ti], + } + + +def plot_efit_qmap(qmap, ece_R_peak=None, ece_rho_peak=None, + q_contours=(1.5, 2.0, 2.5, 3.0), + shot=None, figsize=(6, 7)): + """ + Plot 2-D q map from EFIT with ECE mode location overlaid. + + Filled contours of q on (R, Z), LCFS boundary, magnetic axis. + ECE midplane sightline (Z=0) and ECE peak R marked. + + Parameters + ---------- + qmap : dict returned by fetch_efit_qmap + ece_R_peak : float — major radius of ECE amplitude peak [m] + ece_rho_peak : float — normalised rho for annotation (optional) + q_contours : iterable of q values to draw as labelled contour lines + """ + R = qmap['R_grid'] + Z = qmap['Z_grid'] + q2d = qmap['q_2d'] + t = qmap['t_actual'] + + fig, ax = plt.subplots(figsize=figsize) + title = (f'Shot {shot} — ' if shot else '') + f'q map t={t:.0f} ms (EFIT02ER)' + ax.set_title(title, fontsize=9) + + # Filled q contour (clipped for colour clarity) + q_plot = np.clip(q2d, 0.5, 5.0) + cf = ax.contourf(R, Z, q_plot, levels=40, cmap='RdYlBu_r', alpha=0.85) + cbar = fig.colorbar(cf, ax=ax, pad=0.02) + cbar.set_label('q', fontsize=9) + + # Named q contours + cs = ax.contour(R, Z, q2d, levels=list(q_contours), + colors='k', linewidths=[2.0 if q == 2.0 else 0.8 + for q in q_contours]) + ax.clabel(cs, fmt='q=%.1f', fontsize=7, inline=True) + + # LCFS boundary + ax.plot(qmap['R_bdy'], qmap['Z_bdy'], 'k-', lw=1.5, label='LCFS') + + # Magnetic axis + ax.plot(qmap['R_axis'], qmap['Z_axis'], '+', color='k', ms=10, mew=2) + + # ECE midplane sightline + ax.axhline(0.0, color='royalblue', ls='--', lw=1.2, label='ECE sightline (Z=0)') + + # ECE peak radius + if ece_R_peak is not None: + q_at_peak = float(np.interp(ece_R_peak, R, + q2d[int(np.argmin(np.abs(Z - 0.0))), :])) + lbl = f'ECE peak R={ece_R_peak:.3f} m q={q_at_peak:.2f}' + if ece_rho_peak is not None: + lbl += f' ρ={ece_rho_peak:.2f}' + ax.axvline(ece_R_peak, color='r', ls='-', lw=1.5, label=lbl) + ax.plot(ece_R_peak, 0.0, 'r*', ms=12, zorder=6) + print(f' q at ECE peak (R={ece_R_peak:.3f} m, Z=0): {q_at_peak:.3f}') + + ax.set_xlabel('R (m)') + ax.set_ylabel('Z (m)') + ax.set_aspect('equal') + ax.legend(fontsize=7, loc='upper right') + ax.grid(True, alpha=0.2) + fig.tight_layout() + return fig + + +# ── Helpers ──────────────────────────────────────────────────────────────────── + +def _freq_smooth(x, nsmooth): + """Uniform running average over nsmooth bins along last axis.""" + if nsmooth < 2: + return x + from scipy.ndimage import uniform_filter1d + return uniform_filter1d(x.real, size=nsmooth, axis=-1) + \ + 1j * uniform_filter1d(x.imag, size=nsmooth, axis=-1) \ + if np.iscomplexobj(x) else \ + uniform_filter1d(x, size=nsmooth, axis=-1) + + +def _c95(nsmooth): + """95% coherence confidence level (Fisher z-transform, IDL formula).""" + if nsmooth < 2: + return 0.0 + z = 1.96 / np.sqrt(2.0 * nsmooth - 2.0) + tanh_z = np.tanh(z) + return float(tanh_z ** 2) + + +# ── Core mode analysis ───────────────────────────────────────────────────────── + +def mode_spectrogram(signals, t_ms, phi_deg, + dt_window_ms=4.0, overlap_frac=0.75, + f_min_khz=5.0, f_max_khz=200.0, + f_smooth_khz=2.0, + n_range=(-5, 5), + remove_dc=True): + """ + Compute power spectrogram and toroidal mode number spectrogram. + + Algorithm: + 1. Divide each probe signal into overlapping Hann-windowed time slices. + 2. FFT each slice; optionally smooth spectra in frequency (f_smooth_khz). + 3. Matched-filter mode number fit at each (time, freq) bin: + A_n = |Σ_j F_j · exp(-i·n·φ_j)| / N_probes + 4. Dominant n = argmax(A_n); coherence = A_dom / ΣA_all. + 5. Compute 95% coherence confidence level (c95) from nsmooth. + + Parameters + ---------- + signals : (n_probes, n_t) + t_ms : (n_t,) time in ms + phi_deg : (n_probes,) toroidal angles in degrees + dt_window_ms : FFT window length [ms] + overlap_frac : fractional window overlap (0–1) + f_min_khz, f_max_khz : analysis frequency band [kHz] + f_smooth_khz : frequency smoothing bandwidth [kHz]; 0 disables smoothing + n_range : (n_min, n_max) toroidal mode numbers to test + remove_dc : subtract window mean before FFT + + Returns + ------- + dict with keys: + t_win_ms : (n_win,) window centre times [ms] + freq_khz : (n_freq,) frequency axis [kHz] + power : (n_win, n_freq) total power [G²/kHz] + n_dominant : (n_win, n_freq) dominant toroidal mode number + coherence : (n_win, n_freq) matched-filter coherence (0–1) + mode_amp : dict {n: (n_win, n_freq)} mode amplitude [G/√kHz] + rms_vs_time: dict {n: (n_win,)} RMS per mode [G] + c95 : float 95% coherence confidence level + nsmooth : int frequency bins smoothed + f_min_khz, f_max_khz, n_range, phi_deg, nwin, fs_khz : metadata + """ + n_probes, n_t = signals.shape + dt_ms = float(np.mean(np.diff(t_ms))) + fs_khz = 1.0 / dt_ms + + nwin = int(round(dt_window_ms / dt_ms)) + nstep = max(1, int(round(nwin * (1 - overlap_frac)))) + phi_rad = np.deg2rad(phi_deg) + + freqs_khz = np.fft.rfftfreq(nwin, d=dt_ms) + df_khz = freqs_khz[1] if len(freqs_khz) > 1 else 1.0 + f_mask = (freqs_khz >= f_min_khz) & (freqs_khz <= f_max_khz) + freqs_out = freqs_khz[f_mask] + + nsmooth = max(1, int(round(f_smooth_khz / df_khz))) if f_smooth_khz > 0 else 1 + c95_val = _c95(nsmooth) + + win = np.hanning(nwin) + n_modes = np.arange(n_range[0], n_range[1] + 1) + + # Matched-filter steering vectors: (n_modes, n_probes) + phase_vectors = np.exp(-1j * np.outer(n_modes, phi_rad)) + + starts = np.arange(0, n_t - nwin + 1, nstep) + n_win = len(starts) + t_win_ms = t_ms[starts + nwin // 2] + + n_freq = freqs_out.size + power = np.zeros((n_win, n_freq)) + n_dominant = np.zeros((n_win, n_freq), dtype=int) + coherence = np.zeros((n_win, n_freq)) + mode_amp = {int(n): np.zeros((n_win, n_freq)) for n in n_modes} + + for iw, i0 in enumerate(starts): + seg = signals[:, i0:i0 + nwin].copy() + if remove_dc: + seg -= seg.mean(axis=1, keepdims=True) + + # FFT: (n_probes, n_rfft) + F_full = np.fft.rfft(seg * win[None, :], axis=1) + + # Smooth auto-power in frequency, then extract band + auto = np.abs(F_full) ** 2 # (n_probes, n_rfft) + if nsmooth > 1: + from scipy.ndimage import uniform_filter1d + auto = uniform_filter1d(auto, size=nsmooth, axis=1) + + # Power: mean auto-power over probes, normalised to G²/kHz + power[iw] = auto.mean(axis=0)[f_mask] / fs_khz + + # Matched-filter amplitude: smooth F in frequency before filtering + F = F_full[:, f_mask] # (n_probes, n_freq) + if nsmooth > 1: + from scipy.ndimage import uniform_filter1d + F = (uniform_filter1d(F.real, size=nsmooth, axis=1) + + 1j * uniform_filter1d(F.imag, size=nsmooth, axis=1)) + + # A[n_modes, n_freq] = |phase_vectors @ F| / n_probes + A = np.abs(phase_vectors @ F) / n_probes + + for j, n in enumerate(n_modes): + mode_amp[int(n)][iw] = A[j] + + best = np.argmax(A, axis=0) + n_dominant[iw] = n_modes[best] + total_amp = A.sum(axis=0) + coherence[iw] = A[best, np.arange(n_freq)] / (total_amp + 1e-30) + + # RMS per mode: integrate mode_amp² over frequency + rms_vs_time = {int(n): np.sqrt(np.trapz(mode_amp[int(n)] ** 2, freqs_out, axis=1)) + for n in n_modes} + + return { + 't_win_ms': t_win_ms, + 't_sig_ms': (float(t_ms[0]), float(t_ms[-1])), + 'freq_khz': freqs_out, + 'power': power, + 'n_dominant': n_dominant, + 'coherence': coherence, + 'mode_amp': mode_amp, + 'rms_vs_time': rms_vs_time, + 'c95': c95_val, + 'nsmooth': nsmooth, + 'f_min_khz': f_min_khz, + 'f_max_khz': f_max_khz, + 'n_range': n_range, + 'phi_deg': phi_deg, + 'nwin': nwin, + 'fs_khz': fs_khz, + } + + +def mode_fit_timeslice(signals, t_ms, phi_deg, t0_ms, f0_khz, + dt_window_ms=4.0, f_smooth_khz=2.0, n_range=(1, 5)): + """ + Coherence-weighted phase fit at a single (t0, f0) point. + + Matches IDL phase_fit.pro: weights phase residuals by 1/sigma where + sigma = sqrt((1/coherence - 1) / nsmooth), then returns the weighted + chi-squared for each candidate mode number n. + + Parameters + ---------- + signals, t_ms, phi_deg : from fetch_mirnov + t0_ms, f0_khz : target time and frequency + dt_window_ms : FFT window length [ms] + f_smooth_khz : frequency smoothing bandwidth [kHz] + n_range : (n_min, n_max) mode numbers to test + + Returns + ------- + dict with keys: phi_deg, power_vs_phi, phase_vs_phi, coherence_vs_phi, + n_modes, chi2_vs_n (weighted), n_best, t0_ms, f0_khz, + fit_curve_phi, fit_curve_phase (for the best n) + """ + dt_ms = float(np.mean(np.diff(t_ms))) + nwin = int(round(dt_window_ms / dt_ms)) + fs_khz = 1.0 / dt_ms + df_khz = 1.0 / (nwin * dt_ms) + nsmooth = max(1, int(round(f_smooth_khz / df_khz))) if f_smooth_khz > 0 else 1 + + i0 = np.searchsorted(t_ms, t0_ms) - nwin // 2 + i0 = max(0, min(i0, len(t_ms) - nwin)) + + seg = signals[:, i0:i0 + nwin].copy() + seg -= seg.mean(axis=1, keepdims=True) + + F_full = np.fft.rfft(seg * np.hanning(nwin)[None, :], axis=1) + freqs_khz = np.fft.rfftfreq(nwin, d=dt_ms) + i_f = np.argmin(np.abs(freqs_khz - f0_khz)) + + # Smooth in frequency around target bin + if nsmooth > 1: + from scipy.ndimage import uniform_filter1d + F_sm = (uniform_filter1d(F_full.real, size=nsmooth, axis=1) + + 1j * uniform_filter1d(F_full.imag, size=nsmooth, axis=1)) + auto_sm = uniform_filter1d(np.abs(F_full) ** 2, size=nsmooth, axis=1) + else: + F_sm = F_full + auto_sm = np.abs(F_full) ** 2 + + C = F_sm[:, i_f] # complex amplitude at f0 for each probe + + # Cross-coherence with probe 0 (reference), matching IDL array_spec + cross_auto = np.abs(F_sm[:, i_f]) ** 2 # auto per probe at f0 + cross_cross = np.abs(F_sm[:, i_f] * np.conj(F_sm[0, i_f])) ** 2 + coh_vs_phi = cross_cross / (cross_auto * cross_auto[0] + 1e-30) + coh_vs_phi = np.clip(coh_vs_phi, 1e-6, 1.0) + + # Phase uncertainty weights (IDL phase_fit.pro formula) + sigma = np.sqrt((1.0 / coh_vs_phi - 1.0) / max(nsmooth, 1)) + sigma = np.maximum(sigma, 1e-3) # floor to avoid zero + + phi_rad = np.deg2rad(phi_deg) + n_modes = np.arange(n_range[0], n_range[1] + 1) + + chi2 = np.zeros(len(n_modes)) + w = 1.0 / sigma + for j, n in enumerate(n_modes): + resid = np.angle(C * np.exp(-1j * n * phi_rad)) # rad, wrapped in [-pi, pi] + # Remove the free global phase offset (circular weighted mean) + phi_0 = np.angle(np.sum(np.exp(1j * resid) * w)) + resid = np.angle(np.exp(1j * (resid - phi_0))) + chi2[j] = float(np.sum(resid ** 2 * w) / (np.sum(w) + 1e-30)) + + n_best = n_modes[np.argmin(chi2)] + + # Best-fit phase line: use the fitted intercept so the line aligns with data + resid_best = np.angle(C * np.exp(-1j * n_best * phi_rad)) + phi_0_best = np.angle(np.sum(np.exp(1j * resid_best) * w)) + phi_fit = np.linspace(phi_deg.min() - 20, phi_deg.max() + 20, 200) + phase_fit = np.degrees(np.angle(np.exp(1j * (n_best * np.deg2rad(phi_fit) + phi_0_best)))) + + return { + 'phi_deg': phi_deg, + 'power_vs_phi': np.abs(C) ** 2, + 'phase_vs_phi': np.angle(C), + 'coherence_vs_phi': coh_vs_phi, + 'n_modes': n_modes, + 'chi2_vs_n': chi2, + 'n_best': int(n_best), + 't0_ms': float(t_ms[i0 + nwin // 2]), + 'f0_khz': float(freqs_khz[i_f]), + 'fit_curve_phi': phi_fit, + 'fit_curve_phase': phase_fit, + 'nsmooth': nsmooth, + } + + +def mode_fit_lsq(C, phi_deg, n_range=(-5, 5)): + """ + Multi-harmonic least-squares fit to probe complex amplitudes. + + Matches IDL slice_fit.pro: fits all harmonics n=nmin..nmax simultaneously + via matrix inversion rather than independently. + + Parameters + ---------- + C : (n_probes,) complex FFT amplitudes at one (t, f) point + phi_deg : (n_probes,) toroidal angles [deg] + n_range : (n_min, n_max) harmonics to fit simultaneously + + Returns + ------- + dict with: + coeffs : dict {n: (cos_coeff, sin_coeff)} for each n + amplitude : dict {n: amplitude} (sqrt(cos² + sin²)) + phase_deg : dict {n: phase [deg]} + fit_phi : (200,) angle axis for fitted curve [deg] + fit_B : (200,) fitted magnetic perturbation (normalised) + n_dominant: int, n with largest amplitude + """ + phi_rad = np.deg2rad(phi_deg) + n_min, n_max = n_range + n_list = list(range(n_min, n_max + 1)) + n_basis = len(n_list) * 2 # cos + sin per harmonic + + # Design matrix A: (n_probes, n_basis) — real-valued fitting of Re(C) + # Using |C| projected onto each harmonic via cos/sin basis + n_probes = len(C) + A = np.zeros((n_probes, n_basis)) + for k, n in enumerate(n_list): + A[:, 2 * k] = np.cos(n * phi_rad) + A[:, 2 * k + 1] = np.sin(n * phi_rad) + + y = np.abs(C) * np.cos(np.angle(C)) # Re(C) + + # Least-squares: coeffs = (AᵀA)⁻¹ Aᵀ y + ATA = A.T @ A + try: + cc = np.linalg.solve(ATA, A.T @ y) + except np.linalg.LinAlgError: + cc = np.linalg.lstsq(A, y, rcond=None)[0] + + coeffs = {} + amplitude = {} + phase_deg = {} + for k, n in enumerate(n_list): + a_cos, a_sin = cc[2 * k], cc[2 * k + 1] + coeffs[n] = (float(a_cos), float(a_sin)) + amplitude[n] = float(np.sqrt(a_cos ** 2 + a_sin ** 2)) + phase_deg[n] = float(np.degrees(np.arctan2(a_sin, a_cos))) + + # Fitted curve + phi_fit = np.linspace(-360, 360, 200) + phi_fit_rad = np.deg2rad(phi_fit) + B_fit = np.zeros(200) + for k, n in enumerate(n_list): + a_cos, a_sin = cc[2 * k], cc[2 * k + 1] + B_fit += a_cos * np.cos(n * phi_fit_rad) + a_sin * np.sin(n * phi_fit_rad) + + n_dom = max(n_list, key=lambda n: amplitude[n]) + + return { + 'coeffs': coeffs, + 'amplitude': amplitude, + 'phase_deg': phase_deg, + 'fit_phi': phi_fit, + 'fit_B': B_fit, + 'n_dominant': int(n_dom), + } + + +# ── SVD / MUSIC mode analysis ────────────────────────────────────────────────── + +def mode_svd_spectrogram(signals, t_ms, phi_deg, + dt_window_ms=4.0, overlap_frac=0.75, + n_avg=8, + f_min_khz=5.0, f_max_khz=200.0, + f_smooth_khz=2.0, + n_range=(-5, 5), + n_src=1, + music=True, + remove_dc=True): + """ + SVD-based mode number spectrogram via cross-spectral matrix (CSM) analysis. + + At each (time, frequency) bin, averages `n_avg` neighbouring FFT snapshots + to form the Hermitian CSM: + + S[j,k] = (1/N) Σ_snapshots F_j · conj(F_k) + + Eigendecompose S = U Λ Uᴴ (Hermitian, so all-real eigenvalues). + The dominant eigenvector u₁ represents the coherent spatial mode with the + most power. Mode number is identified by matched-filter projection of u₁ + onto toroidal steering vectors. + + Optionally computes the MUSIC pseudospectrum using the noise subspace: + + P_MUSIC(n) = 1 / (1 − |aₙᴴ u₁|²) [dB] + + which has a sharper peak at the true mode number than the matched filter. + + Parameters + ---------- + signals : (n_probes, n_t) + t_ms : (n_t,) time in ms + phi_deg : (n_probes,) toroidal angles in degrees + dt_window_ms : FFT window length [ms] + overlap_frac : fractional window overlap (0–1) + n_avg : number of FFT snapshots to average for each CSM estimate + f_min_khz, f_max_khz : analysis frequency band [kHz] + f_smooth_khz : frequency smoothing bandwidth [kHz] applied to each FFT + n_range : (n_min, n_max) toroidal mode numbers + n_src : number of assumed signal sources (noise subspace = n_probes - n_src) + music : if True, compute MUSIC pseudospectrum + remove_dc : subtract window mean before FFT + + Returns + ------- + dict with keys (compatible with plot_modespec): + t_win_ms : (n_win,) window centre times [ms] + freq_khz : (n_freq,) frequency axis [kHz] + power : (n_win, n_freq) total power = trace(CSM)/n_probes [G²/kHz] + n_dominant : (n_win, n_freq) mode number from dominant eigenvector + coherence : (n_win, n_freq) λ₁/Σλ — fractional power in dominant mode + mode_amp : dict {n: (n_win, n_freq)} matched-filter amplitude on u₁ + rms_vs_time : dict {n: (n_win,)} RMS per mode [G] + eigenvalues : (n_win, n_freq, n_probes) all eigenvalues, descending + music_pseudo : (n_win, n_freq, n_modes) MUSIC pseudospectrum [dB] or None + c95, nsmooth, f_min_khz, f_max_khz, n_range, phi_deg, nwin, fs_khz, n_avg + """ + n_probes, n_t = signals.shape + dt_ms = float(np.mean(np.diff(t_ms))) + fs_khz = 1.0 / dt_ms + + nwin = int(round(dt_window_ms / dt_ms)) + nstep = max(1, int(round(nwin * (1 - overlap_frac)))) + phi_rad = np.deg2rad(phi_deg) + + freqs_khz = np.fft.rfftfreq(nwin, d=dt_ms) + df_khz = freqs_khz[1] if len(freqs_khz) > 1 else 1.0 + f_mask = (freqs_khz >= f_min_khz) & (freqs_khz <= f_max_khz) + freqs_out = freqs_khz[f_mask] + n_freq = freqs_out.size + + nsmooth = max(1, int(round(f_smooth_khz / df_khz))) if f_smooth_khz > 0 else 1 + c95_val = _c95(nsmooth * n_avg) # effective smoothing includes snapshot averaging + + win = np.hanning(nwin) + n_modes = np.arange(n_range[0], n_range[1] + 1) + + # Conjugate steering vectors for matched filter: exp(-i·n·φ)/√N + # |aₙᴴ u₁| = |Σ_j exp(-i·n·φ_j)·u₁_j| — maximum when n equals the true mode number + a_norm = np.exp(-1j * np.outer(n_modes, phi_rad)) / np.sqrt(n_probes) + + starts = np.arange(0, n_t - nwin + 1, nstep) + n_win = len(starts) + t_win_ms = t_ms[starts + nwin // 2] + + # ── Step 1: compute all FFT snapshots ───────────────────────────────────── + # F_all: (n_win, n_probes, n_freq) complex + F_all = np.zeros((n_win, n_probes, n_freq), dtype=complex) + for iw, i0 in enumerate(starts): + seg = signals[:, i0:i0 + nwin].copy() + if remove_dc: + seg -= seg.mean(axis=1, keepdims=True) + F_full = np.fft.rfft(seg * win[None, :], axis=1) + if nsmooth > 1: + from scipy.ndimage import uniform_filter1d + F_full = (uniform_filter1d(F_full.real, size=nsmooth, axis=1) + + 1j * uniform_filter1d(F_full.imag, size=nsmooth, axis=1)) + F_all[iw] = F_full[:, f_mask] + + # ── Step 2: SVD / MUSIC at each output time via averaged CSM ────────────── + half_avg = n_avg // 2 + power = np.zeros((n_win, n_freq)) + n_dominant = np.zeros((n_win, n_freq), dtype=int) + coherence = np.zeros((n_win, n_freq)) + eigenvalues= np.zeros((n_win, n_freq, n_probes)) + mode_amp = {int(n): np.zeros((n_win, n_freq)) for n in n_modes} + music_pseudo = np.zeros((n_win, n_freq, len(n_modes))) if music else None + + for iw in range(n_win): + # Snapshot block centred on current window + i_lo = max(0, iw - half_avg) + i_hi = min(n_win, iw + half_avg + 1) + snap = F_all[i_lo:i_hi] # (n_snap, n_probes, n_freq) + n_snap = snap.shape[0] + + # CSM[j,l] = E[F_j · conj(F_l)]: (n_freq, n_probes, n_probes) + S = np.moveaxis(snap, 2, 0) # (n_freq, n_snap, n_probes) + CSM = np.einsum('fkj,fkl->fjl', S, S.conj()) / n_snap + + # Eigendecompose (eigh: ascending order, real eigenvalues) + vals, vecs = np.linalg.eigh(CSM) # (n_freq, n_probes), (n_freq, n_probes, n_probes) + vals = np.flip(vals, axis=-1) # descending + vecs = np.flip(vecs, axis=-1) # dominant first + + eigenvalues[iw] = vals + lambda_sum = vals.sum(axis=-1) # (n_freq,) + power[iw] = np.real(lambda_sum) / (n_probes * fs_khz) + coherence[iw] = np.real(vals[:, 0]) / (np.real(lambda_sum) + 1e-30) + + # Dominant eigenvector: u1 shape (n_freq, n_probes) + u1 = vecs[:, :, 0] + + # Matched filter on u1: A[n_modes, n_freq] + # a_norm: (n_modes, n_probes), u1.T: (n_probes, n_freq) + A = np.abs(a_norm @ u1.T) # (n_modes, n_freq) + for j, n in enumerate(n_modes): + mode_amp[int(n)][iw] = A[j] + best = np.argmax(A, axis=0) + n_dominant[iw] = n_modes[best] + + # MUSIC pseudospectrum + if music: + # Signal subspace: first n_src eigenvectors + # |aₙᴴ U_sig|² = Σ_{knfk', a_norm, U_sig) + sig_power = np.sum(np.abs(proj) ** 2, axis=-1) # (n_modes, n_freq) + denom = np.clip(1.0 - sig_power, 1e-10, None) + music_db = 10.0 * np.log10(1.0 / denom) + music_pseudo[iw] = music_db.T # (n_freq, n_modes) + + # ── RMS per mode vs time ────────────────────────────────────────────────── + rms_vs_time = {int(n): np.sqrt(np.trapz(mode_amp[int(n)] ** 2, freqs_out, axis=1)) + for n in n_modes} + + return { + 't_win_ms': t_win_ms, + 't_sig_ms': (float(t_ms[0]), float(t_ms[-1])), + 'freq_khz': freqs_out, + 'power': power, + 'n_dominant': n_dominant, + 'coherence': coherence, + 'mode_amp': mode_amp, + 'rms_vs_time': rms_vs_time, + 'eigenvalues': eigenvalues, + 'music_pseudo': music_pseudo, + 'c95': c95_val, + 'nsmooth': nsmooth, + 'n_avg': n_avg, + 'n_src': n_src, + 'f_min_khz': f_min_khz, + 'f_max_khz': f_max_khz, + 'n_range': n_range, + 'phi_deg': phi_deg, + 'nwin': nwin, + 'fs_khz': fs_khz, + } + + +def plot_svd(result, shot=None, n1rms_dict=None, + coh_thresh=None, onset_ms=None, figsize=(13, 12)): + """ + Five-panel SVD mode analysis plot. + + Panel 1: Power spectrogram (trace of CSM) [dB] + Panel 2: Coherence λ₁/Σλ with c95 threshold line + Panel 3: Dominant mode number (masked by coherence) + Panel 4: MUSIC pseudospectrum at each time (summed over frequency band) + — shows which mode number dominates vs time [dB] + Panel 5: N1RMS (optional) + """ + t = result['t_win_ms'] + f = result['freq_khz'] + pw = result['power'] + coh = result['coherence'] + nd = result['n_dominant'] + rms = result['rms_vs_time'] + c95 = result.get('c95', 0.0) + n_range = result['n_range'] + n_lo, n_hi = n_range + n_modes_all = np.arange(n_lo, n_hi + 1) + music_pseudo = result.get('music_pseudo') # (n_win, n_freq, n_modes) or None + + thresh = coh_thresh if coh_thresh is not None else max(c95, 0.3) + + has_n1rms = n1rms_dict is not None + has_music = music_pseudo is not None + + n_panels = 4 + int(has_n1rms) + if has_music: + n_panels += 1 + ratios = [2.5, 1.5, 2, 2] + ([2] if has_music else []) + ([1] if has_n1rms else []) + fig, axes = plt.subplots(n_panels, 1, figsize=figsize, + gridspec_kw={'height_ratios': ratios}, + sharex=True) + fig.subplots_adjust(hspace=0.06) + ax_iter = iter(axes) + + def _vline(ax): + if onset_ms is not None: + ax.axvline(onset_ms, color='lime', ls='--', lw=0.9, alpha=0.8) + + # ── Power spectrogram ───────────────────────────────────────────────────── + ax = next(ax_iter) + pw_db = 10 * np.log10(pw.T + 1e-20) + vmax = np.percentile(pw_db, 99) + im = ax.pcolormesh(t, f, pw_db, cmap='inferno', + vmin=vmax - 40, vmax=vmax, shading='nearest') + ax.set_ylabel('Freq (kHz)') + ax.set_ylim(result['f_min_khz'], result['f_max_khz']) + plt.colorbar(im, ax=ax, pad=0.01).set_label('Power (dB)', fontsize=7) + title = f'Shot {shot} — SVD Mode Analysis' if shot else 'SVD Mode Analysis' + ax.set_title(title, fontsize=10) + _vline(ax) + + # ── Coherence λ₁/Σλ ───────────────────────────────────────────────────── + ax = next(ax_iter) + im2 = ax.pcolormesh(t, f, coh.T, cmap='viridis', vmin=0, vmax=1, shading='nearest') + ax.axhline(0, color='w', lw=0) # dummy + ax.set_ylabel('Freq (kHz)') + ax.set_ylim(result['f_min_khz'], result['f_max_khz']) + cb2 = plt.colorbar(im2, ax=ax, pad=0.01) + cb2.set_label('λ₁/Σλ', fontsize=7) + n_avg = result.get('n_avg', '?') + ax.text(0.01, 0.97, + f'c95 = {c95:.2f} (n_avg={n_avg}, nsmooth={result.get("nsmooth",1)})', + transform=ax.transAxes, fontsize=7, va='top', color='white') + _vline(ax) + + # ── Dominant mode number ───────────────────────────────────────────────── + ax = next(ax_iter) + nd_masked = np.where(coh.T > thresh, nd.T.astype(float), np.nan) + cmap_n = plt.get_cmap('RdBu_r', n_hi - n_lo + 1) + im3 = ax.pcolormesh(t, f, nd_masked, cmap=cmap_n, + vmin=n_lo - 0.5, vmax=n_hi + 0.5, shading='nearest') + ax.set_ylabel('Freq (kHz)') + ax.set_ylim(result['f_min_khz'], result['f_max_khz']) + cb3 = plt.colorbar(im3, ax=ax, pad=0.01) + cb3.set_label('n', fontsize=7) + cb3.set_ticks(n_modes_all) + _vline(ax) + + # ── RMS per mode vs time ───────────────────────────────────────────────── + ax = next(ax_iter) + for n in n_modes_all: + if n == 0: + continue + lw = 1.4 if abs(n) == 1 else 0.8 + ax.plot(t, rms.get(int(n), np.zeros_like(t)), color=_n_color(n), + lw=lw, label=f'n={n}') + ax.set_ylabel('RMS (G)', fontsize=9) + ax.set_ylim(bottom=0) + ax.legend(loc='upper right', fontsize=6, ncol=4, framealpha=0.5) + ax.grid(True, alpha=0.3, lw=0.5) + _vline(ax) + + # ── MUSIC pseudospectrum vs time ───────────────────────────────────────── + if has_music: + ax = next(ax_iter) + music_sum = music_pseudo.sum(axis=1) # (n_win, n_modes) + n_modes_list = list(np.arange(n_lo, n_hi + 1)) + for j, n in enumerate(n_modes_list): + if n == 0: + continue + lw = 1.4 if abs(n) == 1 else 0.7 + ax.plot(t, music_sum[:, j], color=_n_color(n), lw=lw, label=f'n={n}') + ax.set_ylabel('MUSIC (dB·kHz)', fontsize=9) + ax.legend(loc='upper right', fontsize=6, ncol=4, framealpha=0.5) + ax.grid(True, alpha=0.3, lw=0.5) + _vline(ax) + + # ── N1RMS ──────────────────────────────────────────────────────────────── + if has_n1rms: + ax = next(ax_iter) + t_n = np.array(n1rms_dict.get('time', [])) + d_n = np.abs(np.array(n1rms_dict.get('data', []))) + if t_n.size: + m_n = (t_n >= t[0]) & (t_n <= t[-1]) + t_n, d_n = t_n[m_n], d_n[m_n] + ax.plot(t_n, d_n, 'k', lw=0.7) + ax.axhline(12, color='r', ls='--', lw=0.8) + ax.set_ylabel('N1RMS (G)', fontsize=9) + ax.set_ylim(0, max(15, d_n.max() * 1.1) if d_n.size else 20) + ax.grid(True, alpha=0.3, lw=0.5) + _vline(ax) + + for ax in axes: + ax.set_xlim(t[0], t[-1]) + axes[-1].set_xlabel('Time (ms)') + fig.align_ylabels(axes) + + # Colorbars on spectrogram panels shrink those axes horizontally; align all + # panel right edges so the time axis is physically the same width everywhere. + fig.canvas.draw() + min_x1 = min(ax.get_position().x1 for ax in axes) + for ax in axes: + p = ax.get_position() + ax.set_position([p.x0, p.y0, min_x1 - p.x0, p.height]) + + return fig + + +# ── Plotting ─────────────────────────────────────────────────────────────────── + +# Fixed per-n colors for line plots — distinct and visible on white background. +# n=1 red (primary NTM), n=2 blue, n=3 dark green (was invisible in RdBu_r), +# n=4 orange; negative-n in cooler/muted tones. +_N_COLORS = { + 1: '#d62728', # red + 2: '#1f77b4', # blue + 3: '#2ca02c', # dark green + 4: '#ff7f0e', # orange + -1: '#e377c2', # pink + -2: '#17becf', # cyan + -3: '#8c564b', # brown + 0: '#7f7f7f', # grey +} +def _n_color(n): + return _N_COLORS.get(int(n), plt.cm.tab10(abs(int(n)) % 10)) + +def plot_modespec(result, shot=None, n1rms_dict=None, + coh_thresh=None, onset_ms=None, figsize=(12, 10), + mode_label='n'): + """ + Four-panel modespec-style plot. + + Panel 1: Power spectrogram [dB] + Panel 2: Dominant mode number (masked by coherence > c95 or coh_thresh) + Panel 3: RMS amplitude per mode vs time [G] + Panel 4: N1RMS signal (optional) + + Parameters + ---------- + result : dict from mode_spectrogram + shot : int, for title + n1rms_dict : dict with 'time' and 'data' keys, overlaid on panel 4 + coh_thresh : float override; defaults to result['c95'] (95% confidence) + onset_ms : float, draws vertical dashed line at TM onset + mode_label : 'n' for toroidal (default) or 'm' for poloidal analysis + """ + t = result['t_win_ms'] + f = result['freq_khz'] + pw = result['power'] + nd = result['n_dominant'] + coh = result['coherence'] + rms = result['rms_vs_time'] + n_range = result['n_range'] + c95 = result.get('c95', 0.0) + + thresh = coh_thresh if coh_thresh is not None else max(c95, 0.3) + + n_lo, n_hi = n_range + n_modes_all = np.arange(n_lo, n_hi + 1) + + n_panels = 4 if n1rms_dict is not None else 3 + ratios = [3, 2, 2] + ([1] if n_panels == 4 else []) + fig, axes = plt.subplots(n_panels, 1, figsize=figsize, + gridspec_kw={'height_ratios': ratios}, + sharex=True) + fig.subplots_adjust(hspace=0.06) + + def _vline(ax): + if onset_ms is not None: + ax.axvline(onset_ms, color='lime', ls='--', lw=0.9, alpha=0.8) + + # ── Panel 1: Power spectrogram ──────────────────────────────────────────── + ax = axes[0] + pw_db = 10 * np.log10(pw.T + 1e-20) + vmax = np.percentile(pw_db, 99) + im = ax.pcolormesh(t, f, pw_db, cmap='inferno', + vmin=vmax - 40, vmax=vmax, shading='nearest') + ax.set_ylabel('Frequency (kHz)') + ax.set_ylim(result['f_min_khz'], result['f_max_khz']) + cb = plt.colorbar(im, ax=ax, pad=0.01) + cb.set_label('Power (dB)', fontsize=8) + array_label = 'Toroidal' if mode_label == 'n' else 'Poloidal' + title = f'Shot {shot} — {array_label} Mode Analysis' if shot else f'{array_label} Mode Analysis' + ax.set_title(title, fontsize=10) + _vline(ax) + + # ── Panel 2: Dominant mode number ───────────────────────────────────────── + ax = axes[1] + nd_masked = np.where(coh.T > thresh, nd.T.astype(float), np.nan) + cmap_n = plt.get_cmap('RdBu_r', n_hi - n_lo + 1) + im2 = ax.pcolormesh(t, f, nd_masked, cmap=cmap_n, + vmin=n_lo - 0.5, vmax=n_hi + 0.5, shading='nearest') + ax.set_ylabel('Frequency (kHz)') + ax.set_ylim(result['f_min_khz'], result['f_max_khz']) + cb2 = plt.colorbar(im2, ax=ax, pad=0.01) + cb2.set_label(mode_label, fontsize=8) + cb2.set_ticks(n_modes_all) + nsmooth = result.get('nsmooth', 1) + ax.text(0.01, 0.97, + f'coh > {thresh:.2f} (c95={c95:.2f}, nsmooth={nsmooth})', + transform=ax.transAxes, fontsize=7, va='top', color='white') + _vline(ax) + + # ── Panel 3: RMS per mode vs time ───────────────────────────────────────── + ax = axes[2] + for n in n_modes_all: + if n == 0: + continue + lw = 1.4 if abs(n) == 1 else 0.8 + ax.plot(t, rms.get(int(n), np.zeros_like(t)), color=_n_color(n), + lw=lw, label=f'{mode_label}={n}') + ax.set_ylabel('RMS (G)', fontsize=9) + ax.set_ylim(bottom=0) + ax.legend(loc='upper right', fontsize=6, ncol=4, framealpha=0.5) + ax.grid(True, alpha=0.3, lw=0.5) + _vline(ax) + + # ── Panel 4 (optional): N1RMS ───────────────────────────────────────────── + if n1rms_dict is not None: + ax = axes[3] + t_n = np.array(n1rms_dict.get('time', [])) + d_n = np.abs(np.array(n1rms_dict.get('data', []))) + if t_n.size: + m_n = (t_n >= t[0]) & (t_n <= t[-1]) + t_n, d_n = t_n[m_n], d_n[m_n] + ax.plot(t_n, d_n, 'k', lw=0.7) + ax.axhline(12, color='r', ls='--', lw=0.8, label='12 G') + ax.set_ylabel('N1RMS (G)', fontsize=9) + ax.set_ylim(0, max(15, d_n.max() * 1.1) if d_n.size else 20) + ax.grid(True, alpha=0.3, lw=0.5) + _vline(ax) + + for ax in axes: + ax.set_xlim(t[0], t[-1]) + axes[-1].set_xlabel('Time (ms)') + fig.align_ylabels(axes) + + # Colorbars on spectrogram panels shrink those axes horizontally; align all + # panel right edges so the time axis is physically the same width everywhere. + fig.canvas.draw() + min_x1 = min(ax.get_position().x1 for ax in axes) + for ax in axes: + p = ax.get_position() + ax.set_position([p.x0, p.y0, min_x1 - p.x0, p.height]) + + return fig + + +def plot_slice(slice_result, figsize=(11, 4)): + """ + Two-panel plot for a single time-slice mode fit. + + Left: probe amplitude and phase vs toroidal angle, with best-fit n line. + Right: weighted chi² vs mode number n. + """ + phi = slice_result['phi_deg'] + power = slice_result['power_vs_phi'] + phase = np.degrees(slice_result['phase_vs_phi']) + coh = slice_result.get('coherence_vs_phi', np.ones_like(phi)) + n_modes= slice_result['n_modes'] + chi2 = slice_result['chi2_vs_n'] + n_best = slice_result['n_best'] + t0 = slice_result['t0_ms'] + f0 = slice_result['f0_khz'] + phi_fit = slice_result.get('fit_curve_phi', np.linspace(-50, 320, 200)) + phase_fit = slice_result.get('fit_curve_phase', n_best * phi_fit % 360 - 180) + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize) + + # Left: amplitude bars + phase scatter + fit line + ax1b = ax1.twinx() + ax1.bar(phi, np.sqrt(power), width=4, alpha=0.5, color='C0', label='|B| (a.u.)') + sc = ax1b.scatter(phi, phase, c=coh, cmap='viridis', s=40, + vmin=0, vmax=1, zorder=5, label='Phase (color=coh)') + ax1b.plot(phi_fit, phase_fit, 'C3--', lw=1.2, label=f'n={n_best} fit') + ax1b.set_ylim(-200, 200) + ax1.set_xlabel('Toroidal angle (°)') + ax1.set_ylabel('|B| (a.u.)', color='C0') + ax1b.set_ylabel('Phase (°)', color='C1') + ax1.set_title(f't = {t0:.0f} ms, f = {f0:.1f} kHz') + plt.colorbar(sc, ax=ax1b, label='Coherence', pad=0.12) + lines1, labs1 = ax1.get_legend_handles_labels() + lines2, labs2 = ax1b.get_legend_handles_labels() + ax1.legend(lines1 + lines2, labs1 + labs2, fontsize=8, loc='upper left') + + # Right: weighted chi² vs n + ax2.plot(n_modes, chi2, 'o-', lw=1.2) + ax2.axvline(n_best, color='r', ls='--', lw=1, label=f'n={n_best} (best)') + ax2.set_xlabel('Toroidal mode number n') + ax2.set_ylabel('Weighted phase residual (rad²)') + ax2.set_title('Mode number fit quality') + ax2.legend(fontsize=9) + ax2.grid(True, alpha=0.3) + + fig.tight_layout() + return fig diff --git a/src/tokeye/modespec/classic/mpi_coherence.py b/src/tokeye/modespec/classic/mpi_coherence.py new file mode 100644 index 0000000..c48b564 --- /dev/null +++ b/src/tokeye/modespec/classic/mpi_coherence.py @@ -0,0 +1,294 @@ +""" +Cross-coherence between Mirnov probe coils (MPI) at EPM frequencies (8–20 kHz) +during quiet vs active EPM phases. + +MPI coils are PTDATA from the D3D tree. They sample at ~200 kHz and directly +resolve the 8–20 kHz EPM band without the aliasing that affects ECE (5 kHz). + +Key probe pairs: + Phi = 157° (SET A): MPI66M157D (midplane), MPI3U157D, MPI1U157D (upper) + Phi = 322° (SET B): MPI11M322D (midplane), MPI3A322D, MPI1A322D (upper) + Toroidal separation: 165° → phase shift for n=1 should be ±165° + +Strategy: + 1. Cross-coherence between same-poloidal pairs at different Phi → detects n=1 + 2. Coherence magnitude at 8–20 kHz: quiet vs active + 3. Cross-spectral phase: is it consistent with n=1 in both phases? + +Run: /fusion/projects/codes/conda/omega/envs_public/general/bin/python3 mpi_coherence.py +""" +import sys, os, json +sys.path.insert(0, '/home/yasodak/NTM_premptive_control') +sys.path.insert(0, '/home/yasodak') + +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from scipy.signal import coherence, welch, csd + +SHOTS = [199606, 199607] +LONG_EVENT = {199606: (2746, 4056), 199607: (3496, 4801)} + +# Segment definitions: inter-burst quiet gaps vs active burst periods. +# "Quiet" = low-N1RMS gaps embedded WITHIN the intermittent burst period, +# NOT the pre-onset flat-top (which has different plasma conditions). +# +# 199607 (long event 3496–4801 ms): +# quiet_1 ~3400ms: gap between pre-event burst (3340ms) and long-event onset +# N1RMS5 max = 0.7–2.3 G over 3360–3475ms +# active: main burst period, N1RMS5 routinely 12–28 G +# quiet_2 ~4720ms: gap after last major burst at 4700ms +# +# 199606 (long event 2746–4056 ms, bursts continue to ~4900ms): +# Inter-burst quiet gaps are very short (< 40ms), insufficient for 5kHz ECE. +# Use the pre-onset quiet (2000–2740ms) for 199606 as no inter-burst gap +# is long enough to resolve EPM alias frequencies. +SEGS = { + 199606: { + 'quiet': (2000, 2740), # pre-onset (no extended inter-burst quiet available) + 'active': (3460, 3620), # peak burst period — N1RMS5 consistently 15–20 G + }, + 199607: { + 'quiet': (3360, 3475), # inter-burst quiet gap ~3400ms (0.7–2.3 G) + 'active': (3800, 4700), # main burst period — N1RMS5 12–28 G + 'quiet2': (4720, 4860), # second quiet gap ~4800ms (2.5–2.0 G) + }, +} + +# Probe pairs: (phi_A_name, phi_B_name, poloidal_label, phi_sep_deg) +# phi_sep = B_phi - A_phi — used to predict n=1 phase shift +PROBE_PAIRS = [ + ('MPI66M157D', 'MPI11M322D', 'midplane', 165.0), + ('MPI3U157D', 'MPI3A322D', 'upper-3', 165.0), + ('MPI1U157D', 'MPI1A322D', 'upper-1', 165.0), +] + +# EPM frequency band +EPM_LO = 8000 # Hz +EPM_HI = 20000 # Hz +NPERSEG = 4096 # gives ~50 Hz resolution at 200 kHz + +CACHE_DIR = '/home/yasodak/exp' + + +# ── Caching helpers ───────────────────────────────────────────────────────── +def cache_path(shot, signal): + return os.path.join(CACHE_DIR, str(shot), f'ptdata_{signal}_{shot}.json') + + +def load_or_fetch(shot, signal): + """Return (t_ms, y) from cache; fetch from MDSplus if absent.""" + cp = cache_path(shot, signal) + if os.path.exists(cp): + with open(cp) as f: + d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: + t = t * 1e3 + return t, y + + # Not cached — fetch via MDSplus + print(f' Fetching {signal} for shot {shot} from atlas.gat.com ...') + try: + import MDSplus as mds + conn = mds.Connection('atlas.gat.com') + conn.openTree('D3D', shot) + y = np.array(conn.get(f'PTDATA("{signal}", {shot})').data(), dtype=float) + t = np.array(conn.get(f'DIM_OF(PTDATA("{signal}", {shot}))').data(), dtype=float) + conn.closeAllTrees() + except Exception as e: + print(f' FAILED: {e}') + return None, None + + if t.size == 0: + print(f' Empty result for {signal}') + return None, None + + if np.nanmedian(np.abs(t)) < 10: + t = t * 1e3 # s → ms + print(f' Got {len(t)} pts, t={t[0]:.1f}–{t[-1]:.1f} ms, ' + f'fs≈{1000/np.median(np.diff(t)):.0f} Hz') + + os.makedirs(os.path.join(CACHE_DIR, str(shot)), exist_ok=True) + with open(cp, 'w') as f: + json.dump({'t': t.tolist(), 'y': y.tolist()}, f) + print(f' Cached → {cp}') + return t, y + + +def load_n1rms(shot): + cp = f'{CACHE_DIR}/{shot}/mds_mhd_MHD__TOP_MIRNOV_N1RMS5.json' + with open(cp) as f: + d = json.load(f) + t, y = np.array(d['t']), np.array(d['y']) + if np.nanmedian(t) < 10: + t = t * 1e3 + return t, y + + +def get_segment(t, y, tmin, tmax): + mk = (t >= tmin) & (t <= tmax) & np.isfinite(y) + return t[mk], y[mk] + + +# ── Main loop ──────────────────────────────────────────────────────────────── +for shot in SHOTS: + print(f'\n=== Shot {shot} ===') + segs = SEGS[shot] + + # Load all probe signals + probes = {} + signals_needed = set() + for pA, pB, _, _ in PROBE_PAIRS: + signals_needed.add(pA) + signals_needed.add(pB) + + for sig in sorted(signals_needed): + t, y = load_or_fetch(shot, sig) + if t is not None and len(t) > 100: + probes[sig] = (t, y) + fs_est = 1000.0 / np.median(np.diff(t)) + print(f' {sig}: {len(t)} pts, fs≈{fs_est:.0f} Hz') + else: + print(f' {sig}: not available') + + available_pairs = [(pA, pB, lbl, sep) for pA, pB, lbl, sep in PROBE_PAIRS + if pA in probes and pB in probes] + if not available_pairs: + print(' No probe pairs available — skipping') + continue + + # Estimate sampling rate from first available probe + first_sig = list(probes.keys())[0] + t0, _ = probes[first_sig] + FS = 1000.0 / np.median(np.diff(t0)) # Hz + print(f' Sampling rate: {FS:.0f} Hz') + + # Recalculate NPERSEG to give ~50 Hz resolution + nperseg = max(256, int(FS / 50)) # points per segment → ~50 Hz resolution + + # Figure: one row per probe pair, 3 columns (PSD overlay, coherence, phase) + n_pairs = len(available_pairs) + fig, axes = plt.subplots(n_pairs, 3, figsize=(15, 3.5 * n_pairs), + gridspec_kw={'wspace': 0.3, 'hspace': 0.35}) + if n_pairs == 1: + axes = axes[np.newaxis, :] + q2_lbl = f' quiet2={segs["quiet2"]}ms' if 'quiet2' in segs else '' + fig.suptitle(f'Shot {shot}: MPI cross-coherence at EPM band (8–20 kHz)\n' + f'quiet={segs["quiet"]}ms active={segs["active"]}ms{q2_lbl}', fontsize=11) + + # Build phase list: quiet, active, and optionally quiet2 + phase_list = [ + ('quiet', segs['quiet'], 'tab:green', 'quiet ~3400ms'), + ('active', segs['active'], 'tab:red', 'active burst'), + ] + if 'quiet2' in segs: + phase_list.append(('quiet2', segs['quiet2'], 'tab:blue', 'quiet ~4800ms')) + + print(f'\n EPM band coherence summary ({EPM_LO}–{EPM_HI} Hz):') + hdr = ' ' + ' '.join([f'{"coh_"+ph[0]:>12}' for ph in phase_list]) + \ + ' ' + ' '.join([f'{"phase_"+ph[0]+"°":>13}' for ph in phase_list]) + print(f' {"Pair":<25}{hdr}') + + for row, (pA, pB, lbl, phi_sep) in enumerate(available_pairs): + tA, yA = probes[pA] + tB, yB = probes[pB] + + ax_p = axes[row, 0] # PSD + ax_c = axes[row, 1] # coherence + ax_ph = axes[row, 2] # cross-spectral phase + + coh_band = {} + phase_band = {} + + for (phase_key, (tmin, tmax), color, _) in phase_list: + _, yA_seg = get_segment(tA, yA, tmin, tmax) + _, yB_seg = get_segment(tB, yB, tmin, tmax) + n_pts = min(len(yA_seg), len(yB_seg)) + if n_pts < nperseg * 4: + print(f' {lbl} {phase_key}: insufficient data ({n_pts} pts)') + continue + + yA_s = yA_seg[:n_pts] + yB_s = yB_seg[:n_pts] + + # PSD of probe A + f_p, Pxx = welch(yA_s, fs=FS, nperseg=nperseg) + ax_p.semilogy(f_p, Pxx, color=color, lw=0.8, alpha=0.8, + label=f'{phase_key}') + ax_p.set_xlim(0, 30000) + ax_p.axvspan(EPM_LO, EPM_HI, color='yellow', alpha=0.15, zorder=0) + + # Cross-coherence + f_c, coh = coherence(yA_s, yB_s, fs=FS, nperseg=nperseg) + ax_c.plot(f_c, coh, color=color, lw=0.8, alpha=0.8, + label=f'{phase_key}') + ax_c.axvspan(EPM_LO, EPM_HI, color='yellow', alpha=0.15, zorder=0) + ax_c.set_xlim(0, 30000) + ax_c.set_ylim(0, 1) + + # Cross-spectral phase + f_s, Pxy = csd(yA_s, yB_s, fs=FS, nperseg=nperseg) + phase_xy = np.angle(Pxy, deg=True) + # Weight phase by coherence (plot only where coherence > 0.1) + mask_coh = coh > 0.1 + ax_ph.scatter(f_s[mask_coh], phase_xy[mask_coh], s=1, c=color, + alpha=0.4, label=f'{phase_key}') + ax_ph.axvspan(EPM_LO, EPM_HI, color='yellow', alpha=0.15, zorder=0) + ax_ph.set_xlim(0, 30000) + ax_ph.set_ylim(-180, 180) + ax_ph.axhline(phi_sep, color='k', lw=0.7, ls='--', + label=f'n=1 pred +{phi_sep:.0f}°') + ax_ph.axhline(phi_sep - 360, color='k', lw=0.7, ls=':') + + # Mean coherence in EPM band + epm_mask = (f_c >= EPM_LO) & (f_c <= EPM_HI) + coh_band[phase_key] = float(np.mean(coh[epm_mask])) if epm_mask.sum() > 0 else np.nan + # Median phase in EPM band (weighted by coherence²) + if epm_mask.sum() > 0: + w = coh[epm_mask] ** 2 + ph_vals = phase_xy[epm_mask] + if w.sum() > 0: + phase_band[phase_key] = float(np.average(ph_vals, weights=w)) + else: + phase_band[phase_key] = np.nan + else: + phase_band[phase_key] = np.nan + + # Axis labels and titles + ax_p.set_title(f'{pA} ({lbl})', fontsize=8) + ax_p.set_xlabel('f (Hz)', fontsize=7) + ax_p.set_ylabel('PSD (arb²/Hz)', fontsize=7) + ax_p.legend(fontsize=7) + ax_p.tick_params(labelsize=6) + + ax_c.set_title(f'{pA} × {pB}', fontsize=8) + ax_c.set_xlabel('f (Hz)', fontsize=7) + ax_c.set_ylabel('Coherence', fontsize=7) + ax_c.legend(fontsize=7) + ax_c.tick_params(labelsize=6) + ax_c.axhline(2 / nperseg * np.log(20), color='gray', lw=0.7, ls='--', + label='95% sig.') + + ax_ph.set_title(f'Phase {pA}→{pB} (Δφ={phi_sep:.0f}°)', fontsize=8) + ax_ph.set_xlabel('f (Hz)', fontsize=7) + ax_ph.set_ylabel('Cross-spectral phase (°)', fontsize=7) + ax_ph.legend(fontsize=6, markerscale=5) + ax_ph.tick_params(labelsize=6) + + # Print numerical summary for all phases + coh_vals = ' '.join([f'{coh_band.get(ph[0], float("nan")):>12.3f}' + if not np.isnan(coh_band.get(ph[0], float("nan"))) + else f'{"---":>12}' for ph in phase_list]) + phase_vals = ' '.join([f'{phase_band.get(ph[0], float("nan")):>13.1f}' + if not np.isnan(phase_band.get(ph[0], float("nan"))) + else f'{"---":>13}' for ph in phase_list]) + print(f' {f"{pA}×{pB}":<25} {coh_vals} {phase_vals} (n=1 pred: ±{phi_sep:.0f}°)') + + out = f'figures/mpi_coherence_{shot}.png' + fig.savefig(out, dpi=120, bbox_inches='tight') + print(f' Saved {out}') + plt.close(fig) + +print('\nDone.') diff --git a/src/tokeye/modespec/classic/mre_utils.py b/src/tokeye/modespec/classic/mre_utils.py new file mode 100644 index 0000000..bf95ef4 --- /dev/null +++ b/src/tokeye/modespec/classic/mre_utils.py @@ -0,0 +1,214 @@ +""" +mre_utils.py — Modified Rutherford Equation (MRE) helper functions. + +References: + La Haye et al., Phys. Plasmas 9, 2051 (2002) + Sauter et al., Phys. Plasmas 4, 1654 (1997) + Hegna & Callen, Phys. Plasmas 4, 2940 (1997) + +All quantities SI unless noted. +""" + +import numpy as np + +MU0 = 4 * np.pi * 1e-7 # H/m + + +def resistive_time(eta, r_s, prefactor=1.22): + """ + Resistive diffusion time at the rational surface. + + τ_R = μ₀ r_s² / (prefactor * η) + + Parameters + ---------- + eta : float resistivity [Ω·m] + r_s : float minor radius of q=2 surface [m] + prefactor : float numerical factor (1.22 from resistive MHD; sometimes 1.0) + + Returns + ------- + tau_R : float [s] + """ + return MU0 * r_s**2 / (prefactor * eta) + + +def mre_rhs(W, delta_prime, delta_bs=0.0, delta_eccd=0.0, delta_pol=0.0): + """ + RHS of the Modified Rutherford Equation (normalized). + + (τ_R / r_s) * dW/dt = Δ' + Δ'_bs + Δ'_ECCD + Δ'_pol + + Parameters + ---------- + W : float or array island half-width [m] + delta_prime : float classical tearing index Δ' [m⁻¹] + delta_bs : float bootstrap current term [m⁻¹] (positive = destabilizing) + delta_eccd : float ECCD term [m⁻¹] (negative = stabilizing) + delta_pol : float polarization current term [m⁻¹] + + Returns + ------- + float or array : (τ_R / r_s) * dW/dt [m⁻¹] + """ + return delta_prime + delta_bs + delta_eccd + delta_pol + + +def delta_prime_bootstrap(W, beta_pol, r_s, L_q, L_p, rho_i, C_bs=1.0): + """ + Bootstrap current contribution to the MRE (simplified NTM form). + + Δ'_bs ≈ C_bs * (β_pol / r_s) * (r_s / L_q) * (r_s / L_p) + * W / (W² + ρ_i²) + + The ρ_i² term in the denominator suppresses the bootstrap drive below + the ion orbit width (ion polarization threshold). + + Parameters + ---------- + W : float or array island half-width [m] + beta_pol : float poloidal beta + r_s : float rational surface minor radius [m] + L_q : float q-profile scale length: L_q = q / (dq/dr) [m] + L_p : float pressure scale length: L_p = -p / (dp/dr) [m] + rho_i : float ion poloidal Larmor radius [m] + C_bs : float order-unity geometry coefficient + + Returns + ------- + float or array : Δ'_bs [m⁻¹] + """ + W = np.asarray(W, dtype=float) + if r_s == 0.0 or L_q == 0.0 or L_p == 0.0: + return np.zeros_like(W) + W_thresh = float(rho_i) + denom = W**2 + W_thresh**2 + return C_bs * (beta_pol / r_s) * (r_s / L_q) * (r_s / L_p) * np.where(denom > 0, W / denom, 0.0) + + +def delta_prime_eccd(j_eccd, r_s, B_theta_s, q_s, R0, W, sigma_eccd=None): + """ + ECCD stabilization term in the MRE (Hegna–Callen form). + + Δ'_ECCD ≈ -μ₀ R₀ q_s / B_θ(r_s) * / W * f_shape + + where f_shape accounts for the profile width relative to W. For a + Gaussian with σ_eccd >> W this saturates; for σ_eccd < W it scales ~1. + + Parameters + ---------- + j_eccd : float peak ECCD current density [A/m²] + r_s : float rational surface minor radius [m] + B_theta_s : float poloidal field at rational surface [T] + q_s : float safety factor at rational surface (≈ 2) + R0 : float major radius [m] + W : float or array island half-width [m] + sigma_eccd: float or None 1-σ width of ECCD deposition [m]; None → Gaussian ignored + + Returns + ------- + float or array : Δ'_ECCD [m⁻¹] (negative = stabilizing) + """ + W = np.asarray(W) + prefactor = -MU0 * R0 * q_s / B_theta_s + if sigma_eccd is not None and np.isfinite(sigma_eccd) and sigma_eccd > 0: + # Effective island-width averaged drive (Gaussian profile) + f_shape = np.where(W > 0, np.tanh(W / sigma_eccd), 0.0) + return prefactor * j_eccd * sigma_eccd * f_shape / W + else: + return prefactor * j_eccd / W + + +def seed_island_width_from_Brtilde(Br_tilde, r_wall, r_s, B_theta_s, m=2): + """ + Seed island half-width from radial field fluctuation at the wall. + + Cylindrical extrapolation: + Ψ̃(r_s) = |B̃_r(r_wall)| * r_wall / (m-1) * (r_s / r_wall)^m + Island half-width (La Haye convention): + W_seed = 4 √( r_s |Ψ̃(r_s)| / (m B_θ(r_s)) ) + + Parameters + ---------- + Br_tilde : float or array radial field fluctuation at wall [T] + r_wall : float probe / wall minor radius [m] + r_s : float q=m/n surface minor radius [m] + B_theta_s : float poloidal field at rational surface [T] + m : int poloidal mode number + + Returns + ------- + W_seed : float or array seed island half-width [m] + """ + Br_tilde = np.asarray(Br_tilde) + psi_tilde_s = np.abs(Br_tilde) * r_wall / (m - 1) * (r_s / r_wall)**m + W_seed = 4.0 * np.sqrt(r_s * psi_tilde_s / (m * B_theta_s)) + return W_seed + + +def B_theta_at_surface(Ip_A, r_s, kappa=1.8): + """ + Approximate poloidal field at the rational surface (shaped-cylinder model). + + B_θ(r_s) ≈ μ₀ I_p / (2π r_s κ) + + Parameters + ---------- + Ip_A : float plasma current [A] + r_s : float rational surface minor radius [m] + kappa : float elongation (DIII-D typical ≈ 1.8) + + Returns + ------- + float [T] + """ + return MU0 * Ip_A / (2 * np.pi * r_s * kappa) + + +def find_q_surface(qpsi, aminor, q_target=2.0): + """ + Find the minor radius r where q(r) = q_target via linear interpolation. + + Parameters + ---------- + qpsi : 1-D array safety factor profile on uniform rho grid (0 → 1) + aminor : float minor radius [m] + q_target : float target q value (default 2.0) + + Returns + ------- + r_s : float or None minor radius [m] of q=q_target surface + """ + qpsi = np.asarray(qpsi, dtype=float) + rho = np.linspace(0.0, 1.0, len(qpsi)) + crossings = np.where(np.diff(np.sign(qpsi - q_target)))[0] + if len(crossings) == 0: + return None + i = crossings[0] + frac = (q_target - qpsi[i]) / (qpsi[i + 1] - qpsi[i]) + rho_s = rho[i] + frac * (rho[i + 1] - rho[i]) + return rho_s * aminor + + +def gradient_scale_length(profile, rho_grid, aminor, rho_s): + """ + Compute the profile scale length L = -f / (df/dr) at rho_s. + + Parameters + ---------- + profile : 1-D array profile values on rho_grid + rho_grid : 1-D array normalized radial grid (0–1) + aminor : float minor radius [m] + rho_s : float normalized radius of rational surface + + Returns + ------- + L : float scale length [m] (positive) + """ + r = rho_grid * aminor + dfddr = np.gradient(profile, r) + f_at_s = np.interp(rho_s, rho_grid, profile) + dfddr_at_s = np.interp(rho_s, rho_grid, dfddr) + if dfddr_at_s == 0: + return np.inf + return -f_at_s / dfddr_at_s diff --git a/src/tokeye/modespec/deep/README.md b/src/tokeye/modespec/deep/README.md new file mode 100644 index 0000000..ac1cf0b --- /dev/null +++ b/src/tokeye/modespec/deep/README.md @@ -0,0 +1,20 @@ +# modespec/deep — reserved for the next-generation mode-number engine + +Classic modespec (`../classic`) needs a spatial probe array (toroidal Mirnov +set) to fit mode numbers. The "deep" engine aims to recover toroidal mode +numbers from a *single* line-integrated chord (the DIII-D CO2 interferometer) +using physics side channels instead of spatial phase fits: multitaper +spectrograms, calibrated line detection (no free parameters), per-chord track +building, and harmonic-comb families (a rotating island's harmonics at +`k * f0` carry `n = k * n1`). + +That work lives in the sibling project `integratedmode` +(`/scratch/gpfs/nc1514/integratedmode` — see its `CLAUDE.md` and +`docs/specs/2026-07-01-calibrated-detection-and-n-inference.md`). It stays +there until the analysis is validated; this directory only reserves the +integration point. Intended CLI shape once it lands: + + tokeye modespec --engine deep + +(`--engine classic` stays the default; no breaking changes to the classic +config format.) diff --git a/src/tokeye/modespec/deep/__init__.py b/src/tokeye/modespec/deep/__init__.py new file mode 100644 index 0000000..05ca956 --- /dev/null +++ b/src/tokeye/modespec/deep/__init__.py @@ -0,0 +1,9 @@ +"""Placeholder for next-generation ("deep") modespec — see README.md here. + +Planned backend: single-chord toroidal mode-number identification from the +CO2 interferometer (multitaper spectral lines, harmonic-comb families, +calibrated detection statistics), developed in the sibling ``integratedmode`` +project. Once mature it plugs in as ``tokeye modespec --engine deep``. +""" + +from __future__ import annotations diff --git a/tests/test_alfvenspec.py b/tests/test_alfvenspec.py new file mode 100644 index 0000000..b334500 --- /dev/null +++ b/tests/test_alfvenspec.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import csv + +import numpy as np +import torch +import torch.nn as nn + +from tokeye.alfvenspec import detect, detect_windowed, write_detections_csv + + +class _StubRCNN(nn.Module): + """Returns canned torchvision-style detections; records its input.""" + + def __init__(self, n: int = 3, height: int = 8, width: int = 6): + super().__init__() + self.dummy = nn.Parameter(torch.zeros(1)) + self.seen = None + self.canned = { + "boxes": torch.tensor([[0.0, 0.0, 2.0, 2.0]] * n), + "labels": torch.ones(n, dtype=torch.int64), + "scores": torch.tensor([0.9, 0.6, 0.2][:n]), + "masks": torch.zeros(n, 1, height, width), + } + + def forward(self, images): + self.seen = images + return [self.canned] + + +def test_detect_filters_by_score_and_returns_numpy(): + model = _StubRCNN() + spectrogram = np.random.default_rng(0).normal(size=(8, 6)).astype(np.float32) + + result = detect(spectrogram, model, score_min=0.5) + + assert isinstance(result["boxes"], np.ndarray) + assert result["boxes"].shape == (2, 4) # score 0.2 filtered out + assert result["scores"].shape == (2,) + assert result["labels"].shape == (2,) + assert result["masks"].shape == (2, 8, 6) # channel dim squeezed + + +def test_detect_feeds_single_channel_standardized_image(): + model = _StubRCNN() + spectrogram = (np.random.default_rng(1).normal(size=(8, 6)) * 5 + 50).astype( + np.float32 + ) + + detect(spectrogram, model) + + (img,) = model.seen + assert img.shape == (1, 8, 6) + assert abs(float(img.mean())) < 1e-4 # standardized per-sample + # torch .std() is ddof=1 vs numpy's ddof=0 used for standardization + assert abs(float(img.std()) - 1.0) < 2e-2 + + +def test_detect_honors_explicit_mean_std(): + model = _StubRCNN() + spectrogram = np.full((8, 6), 10.0, dtype=np.float32) + + detect(spectrogram, model, mean=8.0, std=2.0) + + (img,) = model.seen + assert torch.allclose(img, torch.ones_like(img)) + + +def test_detect_windowed_offsets_boxes_to_global_columns(): + model = _StubRCNN(n=1) # one detection (score 0.9) per window call + spectrogram = np.zeros((8, 112), dtype=np.float32) + + result = detect_windowed(spectrogram, model, window_cols=40) + + # windows: [0:40], [40:80], [80:112] -> 3 detections + assert result["boxes"].shape == (3, 4) + np.testing.assert_allclose(result["boxes"][:, 0], [0.0, 40.0, 80.0]) + np.testing.assert_allclose(result["boxes"][:, 2], [2.0, 42.0, 82.0]) + assert result["masks"] is None + + +def test_detect_windowed_folds_sliver_into_previous_window(): + model = _StubRCNN(n=1) + spectrogram = np.zeros((8, 100), dtype=np.float32) + + result = detect_windowed(spectrogram, model, window_cols=40) + + # final 20-column sliver folds into [40:100] instead of being dropped + assert result["boxes"].shape == (2, 4) + np.testing.assert_allclose(result["boxes"][:, 0], [0.0, 40.0]) + + +def test_detect_windowed_falls_back_to_single_window(): + model = _StubRCNN(n=1) + spectrogram = np.zeros((8, 30), dtype=np.float32) + + result = detect_windowed(spectrogram, model, window_cols=40) + + assert result["boxes"].shape == (1, 4) + assert result["masks"] is not None # single window keeps masks + + +def test_write_detections_csv(tmp_path): + out = tmp_path / "ae_detections.csv" + detections = { + "boxes": np.array([[1.0, 2.0, 3.0, 4.0]]), + "labels": np.array([1]), + "scores": np.array([0.9]), + "masks": np.zeros((1, 8, 6)), + } + write_detections_csv(out, [("shot1.npy", detections)]) + + with out.open() as fh: + rows = list(csv.DictReader(fh)) + assert rows[0]["input"] == "shot1.npy" + assert rows[0]["detection"] == "0" + assert [rows[0][k] for k in ("x1", "y1", "x2", "y2")] == ["1.0", "2.0", "3.0", "4.0"] + assert rows[0]["score"] == "0.9" diff --git a/tests/test_cli_suite.py b/tests/test_cli_suite.py new file mode 100644 index 0000000..40c8b7d --- /dev/null +++ b/tests/test_cli_suite.py @@ -0,0 +1,46 @@ +"""CLI tests for the mode-analysis suite subcommands.""" + +from __future__ import annotations + +from tokeye.cli import build_parser, main + + +def test_elmspec_defaults(): + args = build_parser().parse_args(["elmspec", "input.npy"]) + assert args.command == "elmspec" + assert args.inputs == ["input.npy"] + assert args.model is None + assert args.output_dir == "tokeye_elms" + assert args.threshold == 0.5 + assert args.activity_min == 0.1 + assert args.min_gap_cols == 3 + assert args.min_duration_cols == 1 + assert args.fs is None + assert args.png is False + + +def test_alfvenspec_defaults(): + args = build_parser().parse_args(["alfvenspec", "input.npy"]) + assert args.command == "alfvenspec" + assert args.model == "ae_tf_maskrcnn" + assert args.output_dir == "tokeye_ae" + assert args.score_min == 0.5 + assert args.mean is None + assert args.std is None + assert args.save_masks is True + + +def test_modesearch_prints_plan_and_exits_0(capsys): + exit_code = main(["modesearch"]) + assert exit_code == 0 + out = capsys.readouterr().out + assert "not implemented yet" in out + assert "crawler" in out + + +def test_elmspec_missing_input_exits_2(tmp_path, capsys): + exit_code = main(["elmspec", str(tmp_path / "nope_*.npy")]) + assert exit_code == 2 + err = capsys.readouterr().err + assert "No input files found" in err + assert "tokeye example" in err diff --git a/tests/test_eigspec.py b/tests/test_eigspec.py new file mode 100644 index 0000000..69836eb --- /dev/null +++ b/tests/test_eigspec.py @@ -0,0 +1,61 @@ +"""Tests for the vendored eigspec (tokeye.eigspec). + +The import-smoke test is the regression test for the vendoring fixes: +upstream had `lambda`-attribute SyntaxErrors and missing typing imports, +so `import eigspec` failed everywhere. +""" + +from __future__ import annotations + +import subprocess +import sys + +import numpy as np + +from tokeye.eigspec.utils.subspace_identification import covariance_driven_ssi + + +def test_package_imports_without_sklearn_in_subprocess(): + code = ( + "import tokeye.eigspec, tokeye.eigspec.utils.data_extraction, " + "tokeye.eigspec.utils.subspace_identification, sys; " + "assert 'sklearn' not in sys.modules; print('ok')" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_vis_modules_import(): + # spectral_plots carried five of the seven `.lambda` syntax errors + import matplotlib as mpl + + mpl.use("Agg") + from tokeye.eigspec.vis import spectral_plots # noqa: F401 + + +def test_covariance_driven_ssi_recovers_pole_frequency(): + # 2-channel lightly damped 5 Hz oscillation sampled at 100 Hz + fs, f0, zeta = 100.0, 5.0, 0.02 + t = np.arange(0, 20.0, 1.0 / fs) + rng = np.random.default_rng(0) + envelope = np.exp(-zeta * 2 * np.pi * f0 * t) + data = np.column_stack( + [ + envelope * np.sin(2 * np.pi * f0 * t), + envelope * np.cos(2 * np.pi * f0 * t), + ] + ) + 0.01 * rng.normal(size=(t.size, 2)) + + result = covariance_driven_ssi(data, [10, 10, 2]) + + poles = np.linalg.eigvals(result.state_matrix) + freqs_hz = np.abs(np.angle(poles)) * fs / (2 * np.pi) + dampings = -np.log(np.abs(poles)) / np.abs(np.angle(poles)) + assert np.any(np.abs(freqs_hz - f0) < 0.2), freqs_hz + assert np.any(np.abs(dampings - zeta) < 0.01), dampings diff --git a/tests/test_elmspec.py b/tests/test_elmspec.py new file mode 100644 index 0000000..cb490f5 --- /dev/null +++ b/tests/test_elmspec.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import csv + +import numpy as np + +from tokeye.elmspec import ( + ElmEvent, + column_activity, + extract_elm_events, + summarize, + write_events_csv, +) + + +def _mask_with_bursts(bursts: list[tuple[int, int]], width: int = 50, height: int = 10): + """Transient-channel mask with full-column bursts over [start, end] cols.""" + mask = np.zeros((height, width), dtype=np.float32) + for start, end in bursts: + mask[:, start : end + 1] = 1.0 + return mask + + +def test_column_activity_is_fraction_of_active_bins(): + mask = np.zeros((10, 4), dtype=np.float32) + mask[:5, 1] = 1.0 # half the bins in column 1 + mask[:, 2] = 1.0 # all bins in column 2 + activity = column_activity(mask, threshold=0.5) + assert activity.shape == (4,) + assert activity[0] == 0.0 + assert activity[1] == 0.5 + assert activity[2] == 1.0 + + +def test_extract_merges_events_across_small_gaps(): + mask = _mask_with_bursts([(10, 12), (15, 17)]) # gap of 2 columns + events = extract_elm_events(mask, min_gap_cols=3) + assert len(events) == 1 + assert events[0].start_col == 10 + assert events[0].end_col == 17 + + +def test_extract_keeps_events_across_large_gaps(): + mask = _mask_with_bursts([(10, 12), (30, 32)]) + events = extract_elm_events(mask, min_gap_cols=3) + assert [(e.start_col, e.end_col) for e in events] == [(10, 12), (30, 32)] + + +def test_extract_drops_short_events(): + mask = _mask_with_bursts([(5, 5), (20, 24)]) + events = extract_elm_events(mask, min_gap_cols=1, min_duration_cols=2) + assert [(e.start_col, e.end_col) for e in events] == [(20, 24)] + + +def test_extract_empty_mask_gives_no_events(): + mask = np.zeros((10, 50), dtype=np.float32) + assert extract_elm_events(mask) == [] + + +def test_extract_records_peak_activity(): + mask = np.zeros((10, 50), dtype=np.float32) + mask[:5, 10:13] = 1.0 # activity 0.5 + events = extract_elm_events(mask, activity_min=0.3) + assert len(events) == 1 + assert events[0].peak_activity == 0.5 + + +def test_summarize_counts_and_frequency(): + events = [ElmEvent(10, 12, 1.0), ElmEvent(30, 32, 1.0)] + # 100 columns at hop=256, fs=200_000 -> 0.128 s total + summary = summarize(events, n_cols=100, hop=256, fs=200_000.0) + assert summary["n_events"] == 2 + assert np.isclose(summary["elm_freq_hz"], 2 / (100 * 256 / 200_000.0)) + assert np.isclose(summary["duty_cycle"], 6 / 100) + + +def test_summarize_without_fs_leaves_frequency_unset(): + summary = summarize([ElmEvent(0, 1, 1.0)], n_cols=10, hop=256, fs=None) + assert summary["n_events"] == 1 + assert summary["elm_freq_hz"] is None + + +def test_write_events_csv(tmp_path): + out = tmp_path / "elm_events.csv" + events = [ElmEvent(10, 12, 0.75)] + write_events_csv(out, [("shot1.npy", events)], hop=256, fs=200_000.0) + + with out.open() as fh: + rows = list(csv.DictReader(fh)) + assert rows[0]["input"] == "shot1.npy" + assert rows[0]["event"] == "0" + assert rows[0]["start_col"] == "10" + assert rows[0]["end_col"] == "12" + assert float(rows[0]["t_start_s"]) == 10 * 256 / 200_000.0 + assert float(rows[0]["peak_activity"]) == 0.75 + + +def test_write_events_csv_without_fs_leaves_times_blank(tmp_path): + out = tmp_path / "elm_events.csv" + write_events_csv(out, [("shot1.npy", [ElmEvent(0, 3, 1.0)])], hop=256, fs=None) + with out.open() as fh: + rows = list(csv.DictReader(fh)) + assert rows[0]["t_start_s"] == "" diff --git a/tests/test_hub.py b/tests/test_hub.py index 0615d8c..f91712e 100644 --- a/tests/test_hub.py +++ b/tests/test_hub.py @@ -6,13 +6,56 @@ import torch import torch.nn as nn -from tokeye.hub import DEFAULT_MODEL, MODEL_REGISTRY, load_model +from tokeye.hub import ( + DEFAULT_MODEL, + DEFAULT_REPO_ID, + MODEL_REGISTRY, + download_model, + load_model, + repo_for, +) def test_default_model_is_registered(): assert DEFAULT_MODEL in MODEL_REGISTRY +def test_default_model_is_first_in_registry(): + # _build_from_state_dict tries specs in insertion order; the U-Net must + # come first so its checkpoints never construct the R-CNN builder. + assert next(iter(MODEL_REGISTRY)) == DEFAULT_MODEL + + +def test_repo_for_resolves_per_model_repo(): + assert repo_for("big_tf_unet") == DEFAULT_REPO_ID + assert repo_for("ae_tf_maskrcnn") == "nc1/ae_tf_maskrcnn" + # Unknown names (e.g. local paths) fall back to the default repo. + assert repo_for("/some/local/model.pt") == DEFAULT_REPO_ID + + +def test_download_model_uses_per_model_repo(monkeypatch): + seen = {} + + def fake_hf_hub_download(repo_id, filename, **kwargs): + seen["repo_id"] = repo_id + seen["filename"] = filename + return "/fake/path.pt" + + monkeypatch.setattr("tokeye.hub.hf_hub_download", fake_hf_hub_download) + + download_model("ae_tf_maskrcnn") + assert seen == { + "repo_id": "nc1/ae_tf_maskrcnn", + "filename": "ae_tf_maskrcnn_251223.pt", + } + + download_model("ae_tf_maskrcnn", repo_id="someone/else") + assert seen["repo_id"] == "someone/else" + + download_model("big_tf_unet") + assert seen["repo_id"] == DEFAULT_REPO_ID + + def test_load_model_from_registry_downloads_and_loads(tmp_path, monkeypatch): spec = MODEL_REGISTRY["big_tf_unet"] weights_path = tmp_path / spec.filename diff --git a/tests/test_modespec_classic.py b/tests/test_modespec_classic.py new file mode 100644 index 0000000..c870b5d --- /dev/null +++ b/tests/test_modespec_classic.py @@ -0,0 +1,123 @@ +"""Tests for the vendored pymodespec (tokeye.modespec.classic). + +Numeric behavior only — MDSplus fetch paths are untestable off the GA +cluster and are deliberately not covered. +""" + +from __future__ import annotations + +import matplotlib as mpl +import numpy as np + +mpl.use("Agg") + +from tokeye.modespec.classic import detect_modes, load_config, mode_spectrogram +from tokeye.modespec.classic.generate_modes import _contiguous_runs, _fill_gaps + +CLASSIC_DIR = "src/tokeye/modespec/classic" + + +def test_mode_spectrogram_recovers_synthetic_n2_mode(): + # 6 toroidal probes (non-uniform angles, like a real Mirnov array — + # uniform spacing would alias n and n±6), one rotating n=2 mode at 50 kHz + n_true, f_khz = 2, 50.0 + fs_khz = 500.0 + t_ms = np.arange(0, 20.0, 1.0 / fs_khz) + phi_deg = np.array([0.0, 40.0, 90.0, 140.0, 200.0, 250.0]) + phi_rad = np.deg2rad(phi_deg) + + rng = np.random.default_rng(0) + signals = np.cos( + 2 * np.pi * f_khz * t_ms[None, :] + n_true * phi_rad[:, None] + ) + 0.01 * rng.normal(size=(6, t_ms.size)) + + # f_smooth=0: the complex frequency smoothing assumes finite-linewidth + # modes; a zero-linewidth synthetic tone has alternating main-lobe bin + # phases (Hann) that uniform smoothing cancels out. + result = mode_spectrogram(signals, t_ms, phi_deg, f_smooth_khz=0.0) + + # at the injected frequency, mid-signal, the fit must pick n=+2 with + # the largest matched-filter amplitude + jf = int(np.argmin(np.abs(result["freq_khz"] - f_khz))) + iw = result["t_win_ms"].size // 2 + assert int(result["n_dominant"][iw, jf]) == n_true + # coherence = A_best / sum(A_n) over 11 tested n; steering-vector + # crosstalk with 6 probes caps a perfect signal well below 1.0 + assert result["coherence"][iw, jf] > 0.2 + amp_true = result["mode_amp"][n_true][iw, jf] + others = [ + result["mode_amp"][n][iw, jf] + for n in range(*result["n_range"]) + if n != n_true + ] + assert amp_true > max(others) + + +def test_contiguous_runs_and_fill_gaps(): + mask = np.array([False, True, True, False, False, True, False]) + assert list(_contiguous_runs(mask)) == [(1, 2), (5, 5)] + + filled = _fill_gaps(mask, max_gap=2) + assert list(_contiguous_runs(filled)) == [(1, 5)] + + # gaps at the edges are never bridged + assert not filled[0] + assert not filled[6] + + +def test_detect_modes_on_canned_result(): + n_win, n_freq = 20, 5 + t = np.arange(n_win, dtype=float) # 1 ms steps + zeros = np.zeros((n_win, n_freq)) + + n_dominant = np.zeros((n_win, n_freq), dtype=int) + coherence = zeros.copy() + amp2 = zeros.copy() + # an n=2 mode living in windows 5..12, frequency bin 3 + n_dominant[5:13, 3] = 2 + coherence[5:13, 3] = 0.9 + amp2[5:13, 3] = 1.5 + + result = { + "t_win_ms": t, + "freq_khz": np.array([10.0, 20.0, 30.0, 40.0, 50.0]), + "n_dominant": n_dominant, + "coherence": coherence, + "mode_amp": {n: (amp2 if n == 2 else zeros) for n in range(-5, 6)}, + "n_range": (-5, 5), + "c95": 0.1, + } + cfg = { + "coherence_min": 0.5, + "amp_min_G": 0.5, + "min_duration_ms": 2.0, + "merge_gap_ms": 2.0, + } + + events = detect_modes(result, cfg) + + assert len(events) == 1 + event = events[0] + assert event["mode_number"] == 2 + assert event["t_start_ms"] == 5.0 + assert event["t_end_ms"] == 12.0 + assert event["peak_freq_khz"] == 40.0 + assert event["peak_amp_G"] == 1.5 + + +def test_load_config_roundtrip_on_vendored_example(): + global_cfg, shot_cfgs = load_config(f"{CLASSIC_DIR}/modes.yaml") + + assert "output_dir" in global_cfg + assert "atlas" in global_cfg + assert shot_cfgs, "vendored modes.yaml must list at least one shot" + for cfg in shot_cfgs: + assert "shot" in cfg + assert "coherence_min" in cfg # defaults merged in + + +def test_import_does_not_require_mdsplus(): + # the imports at the top of this file must succeed without MDSplus + import sys + + assert "MDSplus" not in sys.modules diff --git a/uv.lock b/uv.lock index e2338a1..0c19110 100644 --- a/uv.lock +++ b/uv.lock @@ -2707,7 +2707,7 @@ wheels = [ [[package]] name = "tokeye" -version = "0.11.0" +version = "0.12.0" source = { editable = "." } dependencies = [ { name = "gradio" }, @@ -2716,6 +2716,7 @@ dependencies = [ { name = "numpy" }, { name = "omegaconf" }, { name = "pydantic" }, + { name = "pyyaml" }, { name = "scipy" }, { name = "tables" }, { name = "torch" }, @@ -2726,6 +2727,9 @@ dependencies = [ ] [package.optional-dependencies] +eigspec = [ + { name = "scikit-learn" }, +] train = [ { name = "h5py" }, { name = "ipykernel" }, @@ -2782,7 +2786,9 @@ requires-dist = [ { name = "pandas", marker = "extra == 'train'" }, { name = "pybaselines", marker = "extra == 'train'" }, { name = "pydantic" }, + { name = "pyyaml" }, { name = "scikit-image", marker = "extra == 'train'" }, + { name = "scikit-learn", marker = "extra == 'eigspec'" }, { name = "scikit-learn", marker = "extra == 'train'" }, { name = "scipy" }, { name = "tables" }, @@ -2796,7 +2802,7 @@ requires-dist = [ { name = "tqdm" }, { name = "wavio", marker = "extra == 'train'" }, ] -provides-extras = ["train"] +provides-extras = ["eigspec", "train"] [package.metadata.requires-dev] dev = [