Using the duqmap function
The api provides a map function duqmap which can be used to map functions over data generated by duqtools, or just data that you specify yourself (as in this example)
import xarray as xr
from duqtools.api import ImasHandle, Run, Runs, duqmap
from pydantic_yaml import parse_yaml_raw_as
from imas2xarray import rebase_all_coords
Read data from runs.yaml
This data is later used to instrument duqtools over which runs to apply the function
name = "../example/runs.yaml"
with open(name) as f:
runs = parse_yaml_raw_as(Runs, f)
for run in runs:
run.data_out.relative_location = None
Define a function to be used with duqmap
Functions should have either Run or ImasHandle as the type of the first argument, duqmap will then take care of converting the input to the correct type
def print_and_return_handle(run: Run):
print(run.dirname)
return str(run.dirname)
this function extracts the dirnames from any runs duqmap operates over:
names = duqmap(print_and_return_handle, runs=runs)
A function to extract the electron temperature of each ImasHandle for each output mapped over by duqmap:
def extract_te(handle: ImasHandle):
return handle.get_variables(['t_e', 'rho_tor_norm'])
Retrieve a dataset for each of the runs specified:
datasets = duqmap(extract_te, runs=runs)
datasets
Rebase and convert all the data into a single dataset with the help of xarray
datasets = rebase_all_coords(datasets, datasets[0])
dataset = xr.concat(datasets, dim='run')
dataset['run'] = names
dataset
Plot our results the way we want
dataset['t_e'].isel(time=0).plot.line(x='rho_tor_norm', hue='run')
