Output Setting

Overview

A module for handling the output setting passed in the get_shots_data method. The output setting is used to handle the output of data from DisruptionPy as it is retrieved. This may include collecting all the data from a request and returning it as a list or streaming outputted data to a file as it is retrieved.

This module defines the abstract class OutputSetting that can have subclasses passed as the output_setting argument to the get_shots_data method. It also provides built-in classes and mappings to easily set the output type for common use cases.

Usage

Currently, these are the options that can be passed as the output_setting argument to get_shots_data:

• An instance of a subclass of OutputSetting
• A string identifier in the _output_setting_mappings dictionary:

  _output_setting_mappings: Dict[str, Type[OutputSetting]] = {
      "dataframe": DataFrameOutputSetting,
      "dataset": DatasetOutputSetting,
      "datatree": DataTreeOutputSetting,
      "df": DataFrameOutputSetting,
      "dict": DictOutputSetting,
      "ds": DatasetOutputSetting,
      "dt": DataTreeOutputSetting,
      "pandas": DataFrameOutputSetting,
      "pd": DataFrameOutputSetting,
      "xarray": DatasetOutputSetting,
      "xr": DatasetOutputSetting,
  }
  

• A file path as a string with its suffix mapped to an OutputSetting type in the _file_suffix_to_output_setting dictionary:

  _file_suffix_to_output_setting: Dict[str, Type[OutputSetting]] = {
      ".cdf": DatasetOutputSetting,
      ".csv": DataFrameOutputSetting,
      ".hdf5": DatasetOutputSetting,
      ".h5": DatasetOutputSetting,
      ".nc": DatasetOutputSetting,
      "/": DictOutputSetting,
  }
  

• A Python list of any other output type request option that can be passed as the OutputSetting argument to get_shots_data (all options listed previously). See OutputSettingList for more details.

Built-in Implementations

Handles output settings for retrieving and saving shot data.

This module provides classes and methods to manage various output settings.

disruption_py.settings.output_setting.DataFrameOutputSetting

Bases: DatasetOutputSetting

Outputs data as a DataFrame.

Source code in disruption_py/settings/output_setting.py

class DataFrameOutputSetting(DatasetOutputSetting):
    """
    Outputs data as a DataFrame.
    """

    def concat(self) -> pd.DataFrame:
        """
        Concatenate the resulting DataFrame.

        Returns
        -------
        pd.DataFrame
            The resulting DataFrame.
        """
        if not self.results:
            logger.critical("Nothing to concatenate!")
            return pd.DataFrame()
        df = super().concat().to_dataframe()
        base = ["shot", "time"]
        cols = base + [c for c in sorted(df.columns) if c not in base]
        return df[cols]

concat

concat() -> pd.DataFrame

Concatenate the resulting DataFrame.

RETURNS	DESCRIPTION
DataFrame	The resulting DataFrame.

Source code in disruption_py/settings/output_setting.py

def concat(self) -> pd.DataFrame:
    """
    Concatenate the resulting DataFrame.

    Returns
    -------
    pd.DataFrame
        The resulting DataFrame.
    """
    if not self.results:
        logger.critical("Nothing to concatenate!")
        return pd.DataFrame()
    df = super().concat().to_dataframe()
    base = ["shot", "time"]
    cols = base + [c for c in sorted(df.columns) if c not in base]
    return df[cols]

disruption_py.settings.output_setting.DataTreeOutputSetting

Bases: SingleOutputSetting

Outputs data as a single DataTree.

Source code in disruption_py/settings/output_setting.py

class DataTreeOutputSetting(SingleOutputSetting):
    """
    Outputs data as a single DataTree.
    """

    def concat(self) -> xr.DataTree:
        """
        Concatenate the resulting DataTree.

        Returns
        -------
        xr.DataTree
            The DataTree containing the results, with shots as keys.
        """
        if not self.results:
            logger.critical("Nothing to concatenate!")
            return xr.DataTree()
        return xr.DataTree.from_dict({str(k): v for k, v in self.results.items()})

concat

concat() -> xr.DataTree

Concatenate the resulting DataTree.

RETURNS	DESCRIPTION
DataTree	The DataTree containing the results, with shots as keys.

Source code in disruption_py/settings/output_setting.py

def concat(self) -> xr.DataTree:
    """
    Concatenate the resulting DataTree.

    Returns
    -------
    xr.DataTree
        The DataTree containing the results, with shots as keys.
    """
    if not self.results:
        logger.critical("Nothing to concatenate!")
        return xr.DataTree()
    return xr.DataTree.from_dict({str(k): v for k, v in self.results.items()})

disruption_py.settings.output_setting.DatasetOutputSetting

Bases: SingleOutputSetting

Outputs data as a single Dataset.

Source code in disruption_py/settings/output_setting.py

class DatasetOutputSetting(SingleOutputSetting):
    """
    Outputs data as a single Dataset.
    """

    def concat(self) -> xr.Dataset:
        """
        Concatenate the resulting Dataset.

        Returns
        -------
        xr.Dataset
            The resulting Dataset.
        """
        if not self.results:
            logger.critical("Nothing to concatenate!")
            return xr.Dataset()
        return xr.concat(self.results.values(), dim="idx")

concat

concat() -> xr.Dataset

Concatenate the resulting Dataset.

RETURNS	DESCRIPTION
Dataset	The resulting Dataset.

Source code in disruption_py/settings/output_setting.py

def concat(self) -> xr.Dataset:
    """
    Concatenate the resulting Dataset.

    Returns
    -------
    xr.Dataset
        The resulting Dataset.
    """
    if not self.results:
        logger.critical("Nothing to concatenate!")
        return xr.Dataset()
    return xr.concat(self.results.values(), dim="idx")

disruption_py.settings.output_setting.DictOutputSetting

Bases: OutputSetting

Outputs data as a dictionary of Datasets.

Source code in disruption_py/settings/output_setting.py

class DictOutputSetting(OutputSetting):
    """
    Outputs data as a dictionary of Datasets.
    """

    def __init__(self, path: str | bool = True):
        """
        Initialize empty DictOutputSetting.

        Parameters
        ----------
        path : str | bool, default = True
            The path for writing results to disk.
            If True, a temporary location will be used.
            If False, no results are written to disk.
        """
        self.results: Dict[int, xr.Dataset] = {}
        if path is True:
            path = os.path.join(get_temporary_folder(), "output")
            if os.path.exists(path) or "pytest" in sys.modules:
                path = tempfile.mkdtemp(dir=get_temporary_folder(), prefix="output-")
        self.path = path

    def _output_shot(self, params: OutputSettingParams):
        """
        Store a single result in the dictionary.

        Parameters
        ----------
        params : OutputSettingParams
            The parameters for outputting shot results.
        """
        self.results[params.shot_id] = params.result

    def get_results(self) -> OutputDictType:
        """
        Get the resulting dictionary.

        Returns
        -------
        Dict[int, xr.Dataset]
            The dictionary of results, with shots as keys.
        """
        return self.results

    def to_disk(self) -> str:
        """
        Save all resulting Datasets into a folder.
        """

        if not self.path:
            return ""
        if os.path.exists(self.path):
            if not os.path.isdir(self.path):
                raise FileExistsError(f"Path already exists! {self.path}")
            if os.listdir(self.path):
                logger.warning("Output folder already exists! {path}", path=self.path)
        else:
            os.makedirs(self.path)

        t = time.time()
        for shot, dataset in self.results.items():
            cdf = os.path.join(self.path, f"{shot}.nc")
            logger.trace("Saving result: {cdf}", cdf=cdf)
            dataset.to_netcdf(cdf)
        logger.info(
            "Saved results in {took:.3f} s: {path}",
            took=time.time() - t,
            path=self.path,
        )
        return self.path

__init__

__init__(path: str | bool = True)

PARAMETER	DESCRIPTION
path	The path for writing results to disk. If True, a temporary location will be used. If False, no results are written to disk. TYPE: str | bool DEFAULT: = True

Source code in disruption_py/settings/output_setting.py

def __init__(self, path: str | bool = True):
    """
    Initialize empty DictOutputSetting.

    Parameters
    ----------
    path : str | bool, default = True
        The path for writing results to disk.
        If True, a temporary location will be used.
        If False, no results are written to disk.
    """
    self.results: Dict[int, xr.Dataset] = {}
    if path is True:
        path = os.path.join(get_temporary_folder(), "output")
        if os.path.exists(path) or "pytest" in sys.modules:
            path = tempfile.mkdtemp(dir=get_temporary_folder(), prefix="output-")
    self.path = path

_output_shot

_output_shot(params: OutputSettingParams)

Store a single result in the dictionary.

PARAMETER	DESCRIPTION
params	The parameters for outputting shot results. TYPE: OutputSettingParams

Source code in disruption_py/settings/output_setting.py

def _output_shot(self, params: OutputSettingParams):
    """
    Store a single result in the dictionary.

    Parameters
    ----------
    params : OutputSettingParams
        The parameters for outputting shot results.
    """
    self.results[params.shot_id] = params.result

get_results

get_results() -> OutputDictType

Get the resulting dictionary.

RETURNS	DESCRIPTION
Dict[int, Dataset]	The dictionary of results, with shots as keys.

Source code in disruption_py/settings/output_setting.py

def get_results(self) -> OutputDictType:
    """
    Get the resulting dictionary.

    Returns
    -------
    Dict[int, xr.Dataset]
        The dictionary of results, with shots as keys.
    """
    return self.results

to_disk

to_disk() -> str

Save all resulting Datasets into a folder.

Source code in disruption_py/settings/output_setting.py

def to_disk(self) -> str:
    """
    Save all resulting Datasets into a folder.
    """

    if not self.path:
        return ""
    if os.path.exists(self.path):
        if not os.path.isdir(self.path):
            raise FileExistsError(f"Path already exists! {self.path}")
        if os.listdir(self.path):
            logger.warning("Output folder already exists! {path}", path=self.path)
    else:
        os.makedirs(self.path)

    t = time.time()
    for shot, dataset in self.results.items():
        cdf = os.path.join(self.path, f"{shot}.nc")
        logger.trace("Saving result: {cdf}", cdf=cdf)
        dataset.to_netcdf(cdf)
    logger.info(
        "Saved results in {took:.3f} s: {path}",
        took=time.time() - t,
        path=self.path,
    )
    return self.path

disruption_py.settings.output_setting.OutputSettingList

Bases: OutputSetting

Handles a list of output settings.

Source code in disruption_py/settings/output_setting.py

class OutputSettingList(OutputSetting):
    """
    Handles a list of output settings.
    """

    def __init__(self, output_setting_list: List[OutputSettingType]):
        """
        Initialize OutputSettingList with a list of output settings.

        Parameters
        ----------
        output_setting_list : List[OutputSettingType]
            A list of output settings to handle.
        """
        self.output_setting_list = [
            resolve_output_setting(individual_setting)
            for individual_setting in output_setting_list
        ]

    def _output_shot(self, params: OutputSettingParams):
        """
        Output a single shot for each output setting in the list.

        Parameters
        ----------
        params : OutputSettingParams
            The parameters for outputting shot results.
        """
        _ = [s.output_shot(params) for s in self.output_setting_list]

    def get_results(self) -> List[OutputType]:
        """
        Get results from each output setting in the list.

        Returns
        -------
        List[OutputType]
            A list of results from each output setting.
        """
        return [s.get_results() for s in self.output_setting_list]

    def to_disk(self) -> List[str]:
        """
        Save each OutputSettingList to disk.
        """
        return [s.to_disk() for s in self.output_setting_list]

__init__

__init__(output_setting_list: List[OutputSettingType])

PARAMETER	DESCRIPTION
output_setting_list	A list of output settings to handle. TYPE: List[OutputSettingType]

Source code in disruption_py/settings/output_setting.py

def __init__(self, output_setting_list: List[OutputSettingType]):
    """
    Initialize OutputSettingList with a list of output settings.

    Parameters
    ----------
    output_setting_list : List[OutputSettingType]
        A list of output settings to handle.
    """
    self.output_setting_list = [
        resolve_output_setting(individual_setting)
        for individual_setting in output_setting_list
    ]

_output_shot

_output_shot(params: OutputSettingParams)

Output a single shot for each output setting in the list.

PARAMETER	DESCRIPTION
params	The parameters for outputting shot results. TYPE: OutputSettingParams

Source code in disruption_py/settings/output_setting.py

def _output_shot(self, params: OutputSettingParams):
    """
    Output a single shot for each output setting in the list.

    Parameters
    ----------
    params : OutputSettingParams
        The parameters for outputting shot results.
    """
    _ = [s.output_shot(params) for s in self.output_setting_list]

get_results

get_results() -> List[OutputType]

Get results from each output setting in the list.

RETURNS	DESCRIPTION
List[OutputType]	A list of results from each output setting.

Source code in disruption_py/settings/output_setting.py

def get_results(self) -> List[OutputType]:
    """
    Get results from each output setting in the list.

    Returns
    -------
    List[OutputType]
        A list of results from each output setting.
    """
    return [s.get_results() for s in self.output_setting_list]

to_disk

to_disk() -> List[str]

Save each OutputSettingList to disk.

Source code in disruption_py/settings/output_setting.py

def to_disk(self) -> List[str]:
    """
    Save each OutputSettingList to disk.
    """
    return [s.to_disk() for s in self.output_setting_list]

disruption_py.settings.output_setting.OutputSettingParams dataclass

Parameters for outputting shot results.

ATTRIBUTE	DESCRIPTION
shot_id	Shot ID. TYPE: int
result	Dataset of shot results. TYPE: Dataset
tokamak	The tokamak for which results are being outputted. TYPE: Tokamak

Source code in disruption_py/settings/output_setting.py

@dataclass
class OutputSettingParams:
    """
    Parameters for outputting shot results.

    Attributes
    ----------
    shot_id : int
        Shot ID.
    result : xr.Dataset
        Dataset of shot results.
    tokamak : Tokamak
        The tokamak for which results are being outputted.
    """

    shot_id: int
    result: xr.Dataset
    tokamak: Tokamak

disruption_py.settings.output_setting.SingleOutputSetting

Bases: DictOutputSetting

Abstract class that outputs data as a single object/file.

Source code in disruption_py/settings/output_setting.py

class SingleOutputSetting(DictOutputSetting):
    """
    Abstract class that outputs data as a single object/file.
    """

    def __init__(self, path: str | bool = True):
        """
        Initialize empty SingleOutputSetting.

        Parameters
        ----------
        path : str | bool, default = True
            The path for writing results to disk.
            If True, a unique temporary location will be used.
            If False, no results are written to disk.
        """
        super().__init__(path=False)
        self.result = None
        if path is True:
            ext = ".csv" if "DataFrame" in self.__class__.__name__ else ".nc"
            path = os.path.join(get_temporary_folder(), f"output{ext}")
            if os.path.exists(path) or "pytest" in sys.modules:
                _, path = tempfile.mkstemp(
                    dir=get_temporary_folder(), prefix="output-", suffix=ext
                )
        self.path = path

    @abstractmethod
    def concat(self) -> OutputSingleType:
        """
        Concatenate the resulting object.

        Returns
        -------
        xr.Dataset | xr.DataTree | pd.DataFrame
            The resulting object.
        """

    def get_results(self) -> OutputSingleType:
        """
        Get the resulting object.

        Returns
        -------
        xr.Dataset | xr.DataTree | pd.DataFrame
            The resulting object.
        """
        logger.debug("Concatenating {tot} shots.", tot=len(self.results))
        self.result = self.concat()
        self.results = {}
        return self.result

    def to_disk(self) -> str:
        """
        Save the resulting object into a file.
        """

        if not self.path:
            return ""
        if os.path.exists(self.path) and os.path.getsize(self.path):
            raise FileExistsError(f"File already exists! {self.path}")

        logger.debug(
            "Saving {type}: {path}", type=self.result.__class__.__name__, path=self.path
        )

        t = time.time()
        for method in ["to_netcdf", "to_csv"]:
            if not hasattr(self.result, method):
                continue
            getattr(self.result, method)(self.path)
            break
        else:
            raise NotImplementedError("Could not save object to file.")
        logger.info(
            "Saved {type} in {took:.3f} s: {path}",
            type=self.result.__class__.__name__,
            took=time.time() - t,
            path=self.path,
        )
        return self.path

__init__

__init__(path: str | bool = True)

PARAMETER	DESCRIPTION
path	The path for writing results to disk. If True, a unique temporary location will be used. If False, no results are written to disk. TYPE: str | bool DEFAULT: = True

Source code in disruption_py/settings/output_setting.py

def __init__(self, path: str | bool = True):
    """
    Initialize empty SingleOutputSetting.

    Parameters
    ----------
    path : str | bool, default = True
        The path for writing results to disk.
        If True, a unique temporary location will be used.
        If False, no results are written to disk.
    """
    super().__init__(path=False)
    self.result = None
    if path is True:
        ext = ".csv" if "DataFrame" in self.__class__.__name__ else ".nc"
        path = os.path.join(get_temporary_folder(), f"output{ext}")
        if os.path.exists(path) or "pytest" in sys.modules:
            _, path = tempfile.mkstemp(
                dir=get_temporary_folder(), prefix="output-", suffix=ext
            )
    self.path = path

concat abstractmethod

concat() -> OutputSingleType

Concatenate the resulting object.

RETURNS	DESCRIPTION
Dataset | DataTree | DataFrame	The resulting object.

Source code in disruption_py/settings/output_setting.py

@abstractmethod
def concat(self) -> OutputSingleType:
    """
    Concatenate the resulting object.

    Returns
    -------
    xr.Dataset | xr.DataTree | pd.DataFrame
        The resulting object.
    """

get_results

get_results() -> OutputSingleType

Get the resulting object.

RETURNS	DESCRIPTION
Dataset | DataTree | DataFrame	The resulting object.

Source code in disruption_py/settings/output_setting.py

def get_results(self) -> OutputSingleType:
    """
    Get the resulting object.

    Returns
    -------
    xr.Dataset | xr.DataTree | pd.DataFrame
        The resulting object.
    """
    logger.debug("Concatenating {tot} shots.", tot=len(self.results))
    self.result = self.concat()
    self.results = {}
    return self.result

to_disk

to_disk() -> str

Save the resulting object into a file.

Source code in disruption_py/settings/output_setting.py

def to_disk(self) -> str:
    """
    Save the resulting object into a file.
    """

    if not self.path:
        return ""
    if os.path.exists(self.path) and os.path.getsize(self.path):
        raise FileExistsError(f"File already exists! {self.path}")

    logger.debug(
        "Saving {type}: {path}", type=self.result.__class__.__name__, path=self.path
    )

    t = time.time()
    for method in ["to_netcdf", "to_csv"]:
        if not hasattr(self.result, method):
            continue
        getattr(self.result, method)(self.path)
        break
    else:
        raise NotImplementedError("Could not save object to file.")
    logger.info(
        "Saved {type} in {took:.3f} s: {path}",
        type=self.result.__class__.__name__,
        took=time.time() - t,
        path=self.path,
    )
    return self.path

disruption_py.settings.output_setting.resolve_output_setting

resolve_output_setting(
    output_setting: OutputSettingType,
) -> OutputSetting

Resolve the output setting to an OutputSetting instance.

PARAMETER	DESCRIPTION
output_setting	The output setting to resolve, which can be an instance of OutputSetting, a string, a dictionary, or a list. TYPE: OutputSettingType

RETURNS	DESCRIPTION
OutputSetting	The resolved OutputSetting instance.

Source code in disruption_py/settings/output_setting.py

def resolve_output_setting(
    output_setting: OutputSettingType,
) -> OutputSetting:
    """
    Resolve the output setting to an OutputSetting instance.

    Parameters
    ----------
    output_setting : OutputSettingType
        The output setting to resolve, which can be an instance of OutputSetting,
        a string, a dictionary, or a list.

    Returns
    -------
    OutputSetting
        The resolved OutputSetting instance.
    """
    if isinstance(output_setting, OutputSetting):
        return output_setting

    if isinstance(output_setting, str):
        # check shortcuts
        output_setting_object = _output_setting_mappings.get(output_setting)
        if output_setting_object is not None:
            return output_setting_object()
        # check suffixes
        for suffix, output_setting_type in _file_suffix_to_output_setting.items():
            if output_setting.endswith(suffix):
                return output_setting_type(output_setting)

    if isinstance(output_setting, list):
        return OutputSettingList(output_setting)

    raise ValueError(f"Invalid output processor {output_setting}")

Custom Implementations

Custom implementations of output type requests must inherit from the OutputTypeRequest abstract class, implementing the abstract methods.

Handles output settings for retrieving and saving shot data.

This module provides classes and methods to manage various output settings.

disruption_py.settings.output_setting.OutputSetting

Bases: ABC

OutputSetting abstract class that should be inherited by all output setting classes.

Source code in disruption_py/settings/output_setting.py

class OutputSetting(ABC):
    """
    OutputSetting abstract class that should be inherited by all output setting classes.
    """

    def output_shot(self, params: OutputSettingParams):
        """
        Output a single shot based on the provided parameters.

        Parameters
        ----------
        params : OutputSettingParams
            The parameters for outputting shot results.
        """
        if hasattr(self, "tokamak_overrides"):
            if params.tokamak in self.tokamak_overrides:
                self.tokamak_overrides[params.tokamak](params)
                return
        self._output_shot(params)

    @abstractmethod
    def _output_shot(self, params: OutputSettingParams):
        """
        Abstract method implemented by subclasses to handle data output for a
        single shot.

        Parameters
        ----------
        params : OutputSettingParams
            The parameters for outputting shot results.
        """

    @abstractmethod
    def get_results(self) -> OutputType:
        """
        Return final output after all shots are processed.

        Returns
        -------
        Any
            The final output results.
        """

    @abstractmethod
    def to_disk(self) -> str | List[str]:
        """
        Save final output to disk.
        """

_output_shot abstractmethod

_output_shot(params: OutputSettingParams)

Abstract method implemented by subclasses to handle data output for a single shot.

PARAMETER	DESCRIPTION
params	The parameters for outputting shot results. TYPE: OutputSettingParams

Source code in disruption_py/settings/output_setting.py

@abstractmethod
def _output_shot(self, params: OutputSettingParams):
    """
    Abstract method implemented by subclasses to handle data output for a
    single shot.

    Parameters
    ----------
    params : OutputSettingParams
        The parameters for outputting shot results.
    """

get_results abstractmethod

get_results() -> OutputType

Return final output after all shots are processed.

RETURNS	DESCRIPTION
Any	The final output results.

Source code in disruption_py/settings/output_setting.py

@abstractmethod
def get_results(self) -> OutputType:
    """
    Return final output after all shots are processed.

    Returns
    -------
    Any
        The final output results.
    """

output_shot

output_shot(params: OutputSettingParams)

Output a single shot based on the provided parameters.

PARAMETER	DESCRIPTION
params	The parameters for outputting shot results. TYPE: OutputSettingParams

Source code in disruption_py/settings/output_setting.py

def output_shot(self, params: OutputSettingParams):
    """
    Output a single shot based on the provided parameters.

    Parameters
    ----------
    params : OutputSettingParams
        The parameters for outputting shot results.
    """
    if hasattr(self, "tokamak_overrides"):
        if params.tokamak in self.tokamak_overrides:
            self.tokamak_overrides[params.tokamak](params)
            return
    self._output_shot(params)

to_disk abstractmethod

to_disk() -> str | List[str]

Save final output to disk.

Source code in disruption_py/settings/output_setting.py

@abstractmethod
def to_disk(self) -> str | List[str]:
    """
    Save final output to disk.
    """

disruption_py.settings.output_setting.OutputSettingParams dataclass

Parameters for outputting shot results.

ATTRIBUTE	DESCRIPTION
shot_id	Shot ID. TYPE: int
result	Dataset of shot results. TYPE: Dataset
tokamak	The tokamak for which results are being outputted. TYPE: Tokamak

Source code in disruption_py/settings/output_setting.py

@dataclass
class OutputSettingParams:
    """
    Parameters for outputting shot results.

    Attributes
    ----------
    shot_id : int
        Shot ID.
    result : xr.Dataset
        Dataset of shot results.
    tokamak : Tokamak
        The tokamak for which results are being outputted.
    """

    shot_id: int
    result: xr.Dataset
    tokamak: Tokamak