tokamunch

Python library and CLI (munchi) for generating IMAS IDS mappings using libtokamap. Given a device name and a set of IDS names, tokamunch expands the IDS schema, calls your mapper for every concrete path, and writes results as JSON, HDF5, or NetCDF.

This is mostly vibes and slop for now.

Documentation

Hosted here giving API reference and CLI syntax.

License

MIT License, see license file for more details.

Installation

Available on pypi

pip install tokamunch                   # core (JSON output only)
pip install "tokamunch[imas]"           # adds HDF5 / NetCDF output via imas-python
pip install "tokamunch[dev]"            # adds test + lint tools

Quickstart

# 1. Scaffold a config file for your device.
munchi init-config --output munchi.toml
# Edit munchi.toml: set mapper.device and point to your libtokamap config.

# 2. Preview the paths that will be mapped.
munchi paths --ids magnetics

# 3. Run the mapping and write results.
munchi map --ids magnetics --output results.json

# 4. Inspect the output.
munchi diff results-prev.json results.json

Mapping development workflow

# 1. Generate a mapping template with annotated stubs.
munchi init-mapping --ids magnetics --leaves-only --output magnetics.json

# 2. Edit magnetics.json: replace {"comment": ""} stubs with real mapper paths.

# 3. Map only the paths defined in your file.
munchi map --ids magnetics --mapping magnetics.json --output results.json

# 4. Check your output against a previous run.
munchi diff baseline.json results.json

# 5. Schema updated? Add stubs for new paths without touching existing entries.
munchi update-mapping --ids magnetics --mapping magnetics.json \
    --output magnetics-updated.json

# 6. Incrementally fill in newly mapped paths from an existing result file.
munchi update --input results.json --output results-full.json \
    --ids magnetics --mapping magnetics-updated.json

# 7. Generate shell completions (paste into your shell rc).
munchi completions bash >> ~/.bash_completion
munchi completions zsh  >> ~/.zshrc
munchi completions fish > ~/.config/fish/completions/munchi.fish

Commands

Key flags for munchi map

Configuration

# munchi.toml

[mapper]
device = "mast"

# Option A — point to a libtokamap config file.
config = "config.toml"

# Option B — supply libtokamap config inline.
# [mapper.config_params]
# mapping_directory = "/path/to/mappings"
# schemas_directory = "/path/to/schemas"

[run]
default_shot = 30420
log_level = "WARNING"      # DEBUG | INFO | WARNING | ERROR | CRITICAL
binary_arrays = false      # base64-encode numpy arrays in JSON output
on_imas_error = "fallback-json"  # "fallback-json" | "raise"

[run.concurrency]
mode = "process"   # "serial" | "thread" | "process"
workers = 8

Override any run.* or mapper.* key on the command line without editing the file:

munchi map --ids magnetics --set run.concurrency.mode=thread \
    --set run.concurrency.workers=4 --output results.json

Multi-shot mapping

# Map shots 47125 and 47130; output goes to results_47125.json and results_47130.json.
munchi map --ids magnetics --shots 47125 47130 --output "results_{shot}.json"

# Map every shot from 47100 to 47200.
munchi map --ids magnetics --shot-range 47100 47200 --output "results_{shot}.json"

Library usage

from tokamunch import MappingContext, TokamapInterface
from tokamunch.mapping import collect_mapped_values
from tokamunch.outputs import write_json_file

# Build a context programmatically (no config file required).
ctx = MappingContext.from_config("munchi.toml", shot=47125)

# Run the mapping.
records, summary = collect_mapped_values(ctx, ids_name="magnetics")
write_json_file("results.json", records)
print(summary)

JSON → IDS objects (in-memory augment + write)

from tokamunch.convert import read_json_records, records_to_ids_objects
from tokamunch.write_ids import write_imas_output

records = read_json_records("results.json")
# Augment with manually constructed MappingRecords here if needed.
write_imas_output("output.nc", records=records, force=True)

Comparing results programmatically

from tokamunch.diff import diff_files, render_diff

entries = diff_files("baseline.json", "updated.json", ids_names=["magnetics"])
print(render_diff(entries, "baseline", "updated", show_unchanged=False))

Shell completions

# bash
munchi completions bash >> ~/.bash_completion
# (add `. ~/.bash_completion` to ~/.bashrc if not already sourced)

# zsh
munchi completions zsh >> ~/.zshrc

# fish
munchi completions fish > ~/.config/fish/completions/munchi.fish

Writing a data-source plugin

Tokamunch plugins must implement the libtokamap Python datasource interface: a class with a get(args) method. libtokamap calls get with a dict containing the mapping's DataSourceArgs plus primitive runtime attributes supplied to Mapper.map(...) when the mapping did not already define the same key. In normal tokamunch runs this includes shot, and may include custom runtime args passed when constructing TokamapInterface.

Plugin structure

1. Implement the data source

# my_package/datasource.py
from __future__ import annotations
from typing import Any
import numpy as np

from tokamunch.plugins import DataSourceMetadata

# Optional — declare execution safety and operational hints.
PLUGIN_METADATA = DataSourceMetadata(
    name="my_source",
    display_name="My Data Source",
    version="1.0.0",
    description="Reads signals from My Device.",
    thread_safe=False,   # set True only if your backend is thread-safe
    process_safe=True,
    requires_network=True,
)

class MyDataSource:
    def __init__(self, host: str) -> None:
        self._host = host

    def get(self, args: dict[str, Any]) -> Any:
        """Called by libtokamap for each IDS path."""
        shot = args["shot"]
        signal = args.get("signal_name", "unknown")
        # ... fetch and return a numpy array or scalar ...
        return np.array([1.0, 2.0, 3.0])

def factory(config: dict[str, Any]) -> MyDataSource:
    return MyDataSource(host=config.get("host", "localhost"))

2. Register the entry point

# pyproject.toml
[project.entry-points."tokamunch.data_sources"]
my_source = "my_package.datasource:factory"

3. Declare the plugin in munchi.toml

[[data_sources]]
mapper_name = "my_source"   # name used in libtokamap mapping files
plugin      = "my_source"   # matches the entry-point name
enabled     = true

[data_sources.args]
host = "my-device-server"

Concurrency safety

tokamunch will emit a warning at mapper creation time if your plugin's PLUGIN_METADATA.thread_safe is False and the user has configured run.concurrency.mode = "thread". Use mode = "process" for plugins that are not thread-safe (e.g. those wrapping C extensions with global state).

Development

pip install -e ".[dev]"
pytest                  # run all tests
ruff check src tests    # lint
ruff format src tests   # format
mypy src                # type-check