Skip to content

TokSearch

TokSearch is a Python framework for retrieving and processing fusion experimental data at scale. You define pipelines that pull diagnostic signals, filter and transform the data, and run analyses across thousands of shots, all in ordinary Python.

What It's For

  • Batch analysis across many plasma shots
  • Preprocessing and feature extraction for machine learning
  • Building reproducible, scalable workflows
  • Flexible backend execution (serial, multiprocessing, Ray, Spark)

A pipeline is a list of shots plus a sequence of steps. You fetch signals into each shot's record, map your own functions over it, and drop records with where when they fail a condition. Then you run it on whichever backend suits the job. The same code runs on a laptop or a cluster.

from toksearch import Pipeline, MdsSignal

pipeline = Pipeline(shots)
pipeline.fetch("ip", MdsSignal(r"\ipmhd", "efit01"))
pipeline.map(my_analysis)
results = pipeline.compute_multiprocessing(num_workers=16)

Where the data comes from

TokSearch reads legacy fusion formats directly, so nothing needs converting or copying first:

Source Class Provided by
MDSplus trees MdsSignal toksearch
Zarr stores (object storage) ZarrSignal toksearch
DIII-D PTDATA diagnostics PtDataSignal, RDataSignal toksearch_d3d
DIII-D IMAS IDS paths ImasSignal / D3dImasSignal toksearch_d3d
MAST / MAST-U MastSignal, MastImasSignal toksearch_mast

Device support is a plug-in point: each facility contributes a package of signal classes, and TokSearch core stays device-neutral. The IMAS classes come in device-prefixed pairs for the same reason. D3dImasSignal is an exact alias for ImasSignal, so a script reading IMAS paths from more than one device can name them symmetrically:

from toksearch_d3d import D3dImasSignal
from toksearch_mast import MastImasSignal

Within FDP, those reads go through the federated data layer, so you are not managing file transfers.

Time alignment and datasets

Signals arrive on their own time bases. fetch_dataset and align put several of them on a common one and return an xarray Dataset, which is usually what you want before comparing signals to each other. Results come back as ordinary NumPy, xarray and pandas objects, so the rest of your toolchain works unchanged.

Getting results out

Pipeline.write writes one file per shot in the worker that produced it. That is faster than concatenating on the driver, and it keeps a failed shot from quietly contaminating the output. For smaller results, RecordSet.to_dataframe and to_parquet collect on the driver instead.

A run can also record what produced it. Pass a provenance backend to compute_* and TokSearch captures the shot source, every signal's specification, the operations applied, the backend, and the git commit of the script that ran. See Provenance and CMF.

Documentation

📚 Full documentation, tutorials and API reference: 👉 https://ga-fdp.github.io/toksearch/

The AI assistant writes TokSearch pipelines from plain-language requests if you would rather describe the analysis than code it.

Source Code